leetcode-go/hard/uf.go

package main

// UF ...
type UF struct {
	ID   []int
	Size []int
	Cnt  int
}

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

// Find ...
func (uf UF) Find(id int) int {
	if uf.ID[id] != id {
		uf.ID[id] = uf.Find(uf.ID[id])
	}
	return uf.ID[id]
}

// Union ...
func (uf UF) Union(i, j int) {
	fi, fj := uf.Find(i), uf.Find(j)
	if fi == fj {
		return
	}
	if uf.Size[fi] < uf.Size[fj] { // Swap fi, fj
		fi, fj = fj, fi
	}
	uf.ID[fj] = fi // Merge fj into fi
	uf.Size[fi] += uf.Size[fj]
	uf.Cnt--
}