GCJ/2008/1B/B/uf.go

package main

// UnionFind ...
type UnionFind struct {
	Cnt  int64
	ID   []int64
	Size []int64
}

// NewUF ...
func NewUF(size int64) UnionFind {
	var uf UnionFind
	uf.Cnt = size
	uf.ID = make([]int64, size)
	for i := int64(0); i < size; i++ {
		uf.ID[i] = i
	}
	uf.Size = make([]int64, size)
	return uf
}

// Union ...
func (uf *UnionFind) Union(p, q int64) {
	pID, qID := uf.Find(p), uf.Find(q)
	if pID == qID {
		return
	}
	if uf.Size[pID] < uf.Size[qID] {
		uf.ID[pID] = qID
		uf.Size[qID] += uf.Size[pID]
	} else {
		uf.ID[qID] = pID
		uf.Size[pID] += uf.Size[qID]
	}
	uf.Cnt--
}

// Find ...
func (uf *UnionFind) Find(id int64) int64 {
	for id != uf.ID[id] {
		id = uf.ID[id]
	}
	return id
}