stdUnionFind.go

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

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

func NewWqupc(n int) UnionFinder {
	result := &wqupc{}
	result.ids = make([]int, n)
	result.w = make([]int, n)
	for x := range result.ids {
		result.ids[x] = x
		result.w[x] = 1
	}
	return result
}

func (q *wqupc) Connected(p, r int) bool {
	return q.root(p) == q.root(r)
}

func (q *wqupc) Union(p, r int) {
	qp := q.root(p)
	qr := q.root(r)
	if qr == qp {
		return
	}

	if q.w[qr] > q.w[qp] {
		q.ids[qp] = qr
		q.w[qr] += q.w[qp]
		return
	}
	q.ids[qr] = qp
	q.w[qp] += q.w[qr]
}

func (q *wqupc) root(r int) int {
	for {
		if r == q.ids[r] {
			return r
		}
		q.ids[r] = q.ids[q.ids[r]]
		r = q.ids[r]
	}
}