cleanUnionFind.go

type UnionFinder interface {
   Union(p, r int)
   Connected(p, r int) bool
}

// clean : Weighted Quick Union with Path Compression algorithm
type clean struct {
	ids     []int
	weights []int
}

func NewClean(size int) UnionFinder {
	uf := &clean{}
	uf.initializeArrays(size)
	return uf
}

func (cleanUF *clean) Connected(nodeOne, nodeTwo int) bool {
	return cleanUF.root(nodeOne) == cleanUF.root(nodeTwo)
}

func (cleanUF *clean) Union(nodeOne, nodeTwo int) {
	rootOne := cleanUF.root(nodeOne)
	rootTwo := cleanUF.root(nodetwo)
	if haveTheSameRoot(rootOne, rootTwo) {
		return
	}

	smallTreeRoot, largeTreeRoot := cleanUF.sortRootsByTreeLength(rootOne, rootTwo)
	cleanUF.attachSmallTreeToLargeTreeRoot(smallTreeRoot, largeTreeRoot)
	cleanUF.addSmallTreeLengthToLargeOne(smallTreeRoot, largeTreeRoot)
}

//
// privates
//
func (cleanUF *clean) root(node int) int {
	for {
		if cleanUF.isRoot(node) {
			return node
		}

		node = cleanUF.attachToRootsRoot(node)
	}
}

func (cleanUF *clean) initializeArrays(size int) {
	cleanUF.ids = intVectorOfSize(size)
	cleanUF.weights = intVectorOfSize(size)

	for x := 0; x < size; x++ {
		cleanUF.ids[x] = x
		cleanUF.weights[x] = 1
	}
}

func intVectorOfSize(size int) []int {
	return make([]int, size)
}

func (cleanUF *clean) addSmallTreeLengthToLargeOne(smallTreeRoot int, largeTreeRoot int) {
	cleanUF.weights[largeTreeRoot] += cleanUF.weights[smallTreeRoot]
}

func (cleanUF *clean) attachSmallTreeToLargeTreeRoot(smallTreeRoot, largeTreeRoot int) {
	cleanUF.ids[smallTreeRoot] = largeTreeRoot
}

func (cleanUF *clean) sortRootsByTreeLength(rootOne, rootTwo int) (int, int) {
	if cleanUF.weights[rootOne] < cleanUF.weights[rootTwo] {
		return rootOne, rootTwo
	}

	return rootTwo, rootOne
}

func haveTheSameRoot(rootFirst, rootSecond int) bool {
	return rootFirst == rootSecond
}

func (cleanUF *clean) attachToRootsRoot(node int) int {
	newRoot := cleanUF.ids[cleanUF.ids[node]]
	cleanUF.ids[node] = newRoot
	return newRoot
}

func (cleanUF *clean) isRoot(r int) bool {
	return r == cleanUF.ids[r]
}