golangでUnionFindを作る

はじめに

Union-Findは素集合のデータ構造を表すのにとても便利です.
ここではそのUnion-FindをGo言語で作成していきます.
Union-Find自体には特に触れないのでご注意ください.

どんな機能が必要か

自分的に以下の機能が必要と考えました.

• ある要素aの所属するグループを返す
• ある要素aの所属するグループ内の要素数(木のサイズ)を返す
• ある要素aの所属する[グループA] と ある要素bの所属する[グループB] を合体する
• ある要素aとある要素bが同じグループに所属するかどうかを返す

実装

上の機能をもつように構造体を作成していきましょう.

type UnionFind struct {
	par []int
}

/* コンストラクタ */
func newUnionFind(N int) *UnionFind {
	u := new(UnionFind)
	u.par = make([]int, N)
	for i := range u.par {
		u.par[i] = -1
	}
	return u
}

/* xの所属するグループを返す */
func (u UnionFind) root(x int) int {
	if u.par[x] < 0 {
		return x
	}
	u.par[x] = u.root(u.par[x])
	return u.par[x]
}

/* xの所属するグループ と yの所属するグループ を合体する */
func (u UnionFind) unite(x, y int) {
	xr := u.root(x)
	yr := u.root(y)
	if xr == yr {
		return
	}
	u.par[yr] += u.par[xr]
	u.par[xr] = yr
}

/* xとyが同じグループに所属するかどうかを返す */
func (u UnionFind) same(x, y int) bool {
	if u.root(x) == u.root(y) {
		return true
	}
	return false
}

/* xの所属するグループの木の大きさを返す */
func (u UnionFind) size(x int) int {
	return -u.par[u.root(x)]
}

利用

func main() {
	// インスタンスの生成
	u := newUnionFind(N)
	
	// 各メソッドの利用
	u.unite(a, b)
	u.same(a, b)
	u.size(a)
	u.root(a)
}

内容について

親の要素には, その木の大きさをNとすると, $N \times (-1)$が与えられており,
子の要素には, それが属する木の親の値が与えられるようになっています.

そのため, 要素が負であれば親であることがわかり, それを正の数に戻すことでその木の大きさを得ることできます.