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)$が与えられており,
子の要素には, それが属する木の親の値が与えられるようになっています.

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