AtCoderで使用する言語をGoに変えたので、競プロ用ライブラリをGoで書き直す！
第一弾UnionFind編
最後に全コードをまとめているので、コードだけ見せろという方は記事の下の方まで。

UnionFindとは

素集合を取り扱うデータ構造。要素をグルーピングする時に使う。

集合の代表者を根、要素間の親子関係にあるものをそのまま親、子と呼ぶことにする。

実装

UnionFind構造体のフィールドはシンプル。

type UnionFind struct {
    n    int   // 全要素数
    root []int
}

rootは要素が根かどうかによって表す意味が変わってくる。ここがややこしいポイントだった。要素が根の場合はその集合に含まれる要素数を表し、根でない場合はその要素が含まれる集合の根^[1]はどの要素かを表す。しかも、要素数と根のどちらを表しているかを区別するために、根である場合はその集合の要素数を負の数で表現する^[2]のでこの記事もそうする。
ここさえ理解できればあとは難しくはない。空でも実装できるようになる。

コンストラクタ

func newUnionFind(n int) UnionFind {
	root := make([]int, n)
	for i := 0; i < n; i++ {
		root[i] = -1
	}
	uf := UnionFind{n: n, root: root}
	return uf
}

各要素は初めバラバラの状態なのでそれぞれが自身のみからなる集合の根であり、その集合の要素数は1。つまりrootの初期値は-1になる。(根の時は-(集合の要素数)なので)

find

find(x)はxが属する集合の根を返す。

func (u *UnionFind) find(x int) int {
	if u.root[x] < 0 {
		return x
	}
	u.root[x] = u.find(u.root[x])　// 再帰で根に直接繋ぎ直す
	return u.root[x]
}

xが根の場合
root[x]には負数が入っているため、それで判別する。自身を返して終わり。
xが根でない場合
root[x]にはxの親(根ではない可能性がある)が入っているのでその次の親、次の次の親...と根に至るまで再帰で探索する。
根に至ったら帰りがけで探索してきた要素のrootを根に書き換える。この操作により次回探索時の計算量が落ちる^[3]。この経路圧縮がUnionFindを(計算量的に)使い勝手の良いものにするキモの部分。経路圧縮によりroot[x]には根が入っているので、それを返す。

union

find関数がやっていることと、rootに入っている情報が何かわかっていれば、union関数は簡単。

func (u *UnionFind) union(x, y int) {
	x = u.find(x) // xを根に書き換え
	y = u.find(y) // yを根に書き換え
	if x == y {
        // 同じ根だったら同じ集合に属しているのでunion済み。returnする
		return
	}
	if -u.root[x] < -u.root[y] {
		x, y = y, x
	} // xの集合の要素数の方が大きいようにスワップ
	u.root[x] += u.root[y] // -u.root[x] = (-u.root[x]) + (-u.root[y])をやっている
	u.root[y] = x　// 根xの集合に取り込まれたため、root[y]はxになる
}

集合どうしのマージは、要素数が多い集合に要素数が小さい集合をマージ(by size)しているが、集合の木の高さによるマージ法(by lank)もあるみたい^[4]。rootの表す情報が負数(集合の要素数)だったり根だったりでややこしいが、ちゃんと読めばわかると思う。

おまけの関数

同じ集合に属しているかは根が同じかどうか

func (u *UnionFind) isSame(x, y int) bool {
	return u.find(x) == u.find(y)
}

xが属する集合の要素数

func (u *UnionFind) size(x int) int {
	return -u.root[u.find(x)]
}

まとめると

type UnionFind struct {
	n    int
	root []int
}

func newUnionFind(n int) UnionFind {
	root := make([]int, n)
	for i := 0; i < n; i++ {
		root[i] = -1
	}
	uf := UnionFind{n: n, root: root}
	return uf
}
func (u *UnionFind) find(x int) int {
	if u.root[x] < 0 {
		return x
	}
	u.root[x] = u.find(u.root[x])
	return u.root[x]
}
func (u *UnionFind) union(x, y int) {
	x = u.find(x)
	y = u.find(y)
	if x == y {
		return
	}
	if -u.root[x] < -u.root[y] {
		x, y = y, x
	} // xの方がサイズ大きいように
	u.root[x] += u.root[y]
	u.root[y] = x
}
func (u *UnionFind) isSame(x, y int) bool {
	return u.find(x) == u.find(y)
}
func (u *UnionFind) size(x int) int {
	return -u.root[u.find(x)]
}

参考

書いている途中にUnionFindについての優良記事が出てきた...。ここまで読んでくれた人には申し訳ないがこれを読んでください笑

計算量についてはこちらが詳しく書かれています。

脚注

常に根を格納しているわけではない。親要素の時もある ↩︎
根の要素を負数で表しているのは見たことないかも ↩︎
計算量解析は結構難しそうなので割愛 ↩︎
実装には新たにランク情報を格納するフィールドが必要になる ↩︎

Discussion