KotlinでUnionFind - little star's memory

先日行われたAtCoder Beginner Contest 157 - AtCoderのD問題で、UnionFindを用いる問題が出題されました。

これをきっかけに、KotlinでUnionFindを行うためのコードを整備しました。ご自由にお使いください。

UnionFindとは

UnionFindとは、以下の2種類のクエリを効率的にこなすためのデータ構造です。

unite(x,y) … xの含まれるグループとyの含まれるグループを併合する。
same(x,y) … xとyが同じグループに含まれるか判定する。

はじめはすべてのグループの大きさが1です(すなわちどの2つも同じグループに属しません)。

UnionFindでは、1つのグループを根付き木で表現します。したがって全体では森になります。次のような関数を考えます。

find(x) … xを含む根付き木の根を求める。

これを用いると、same(x,y)は、find(x)とfind(y)が等しいことと言い換えられます。

ということで、unite(x,y)とfind(x)を実装すればよいことになりました。

まずfind(x)を実装しましょう。根付き木なので、根以外には親がただ1つ定まっています。find(x)は、xが根ならxを返します。xが根でないなら、xの親をみて根かどうか調べます。根に到達するまでこれを再帰的に繰り返せばよいです。

次にunite(x,y)を考えます。もともとx,yが同じグループのときは何もしません。そうでないとき、x1,y1をx,yの属する木の根とします(つまり、x1=find(x), y1=find(y)とします)。このとき、x1の親をy1とする、またはy1の親をx1とすることで、2つの根付き木が1つになります。

これでひとまずUnionFindを実装出来ました。

UnionFindの高速化

find(x)の計算では木の高さに比例した時間がかかります。特に縦長になってしまうと、計算量がO(n)になってしまいます。計算量を抑えるためには、木の高さを低くすることが大切です。次の2種類のテクニックがあります。

ランク

unite(x,y)を行う際、2つの木の高さを比較してなるべく高くならないようにつなげます。具体的には、x1=find(x), y1=find(y)とするとき、x1のランクの方が高いときはy1の親をx1に、y1のランクの方が高いときはx1の親をy1にします。こうすることで高さは増加しません。2つの木のランクが同じときは、どちらにしても1だけ高くなります。

経路圧縮

実は木の形はさほど重要ではなく、親さえわかれば十分です。そこで、find(x)をして根を求めたとき、xの親を直接根にしてしまいます。根を求める途中で通った頂点もすべて根につなぎ直します。こうすることで木の高さが抑えられます。なお経路圧縮をすると木の高さが変わることがありますが、ランク計算においてそれは無視することにします。

この2つの工夫により、UnionFindの計算量はO(α(n))となります。α(n)はアッカーマン関数A(n,n)の逆関数で、A(n,n)は非常に大きい関数のため、逆関数であるα(n)は非常に定数に近くなります。

実装例

以下はKotlinによるUnionFind実装例です。

private class UnionFind(n: Int) {
    private val rank = IntArray(n) { 0 }
    private val parent = IntArray(n) { -1 }
 
    fun find(x: Int): Int {
        if (parent[x] < 0) {
            return x
        }
        parent[x] = find(parent[x])
        return parent[x]
    }
 
    fun unite(x: Int, y: Int) {
        val x1 = find(x)
        val y1 = find(y)
        if (x1 == y1) {
            return
        }
        if (rank[x1] < rank[y1]) {
            parent[y1] += parent[x1]
            parent[x1] = y1
        } else {
            parent[x1] += parent[y1]
            parent[y1] = x1
            if (rank[x1] == rank[y1]) {
                rank[x1]++
            }
        }
    }
 
    fun same(x: Int, y: Int): Boolean {
        return find(x) == find(y)
    }
 
    fun size(x: Int): Int {
        return -parent[find(x)]
    }
}

parentは、根以外の場合は自分の親、根の場合は木のサイズの-1倍を表します。

これを用いてABC157Dを解いたものがこちらです。

Submission #10496984 - AtCoder Beginner Contest 157

ACできたので正常に動くはずです。