はじめに

最近、データ構造づいていた(スタック、キュー、ヒープ)けど、今回はUnion-Findをswiftに導入してみる。
ヒープは、溜め込んだデータの内、最大（最小）のデータを取り出す計算量がO(1)と極めて小さく、データを溜め込むときの計算量もO(logN)と比較的小さかった。
Union-Findの特徴は、非常に小さい計算量で、データを「グループ分け」する事が出来ること。
まず、初期化時点で全てのデータ(n個の場合、0..<nのインデックスが付された要素)が存在し、データの追加や削除はされない。メソッドは、unite(同じグループにする)とisSame(同じグループかどうか判定する)の2つのみ。統合（unite）は出来るが、分離は出来ない。また、同じグループかどうかが分かるだけであり、グループ名は得られない。

Union-Findの考え方

uniteとisSameを少ない計算量で実現するのが、Union-Find木であるが、なぜ、treeだと、少ない計算量で計算できるのかは、下記のAtCoderのスライドを見て欲しい。p5のように、配列を使って、「同じ数字なら同じグループ」のようなアルゴリズムだと、2つの要素をuniteしたとき、片方のグループに割り当てた数字を全て塗り返さなくてはならず、O(n)の計算量が必要となり、無駄に計算量を喰ってしまう。

これにたいし、p7のようにtree構造を考え、p9のように「根が同じであれば、同じグループである」とすれば、uniteの計算量は小さくできる。

Union-Find木を使った実装化

c++用のスライドのコードをswiftで実装する

struct UnionFind {
    
    private var par:[Int] = [] // treeの親ノードを格納する。
    
    init(_ n:Int) {
        for i in 0..<n {par.append(i)} 
    }
    
    mutating func root(_ i:Int)->Int {
        if par[i] == i {return i} // 自身が根なら自身を返す
        par[i] = root(par[i])  // 再帰計算の途中で経路圧縮もついでに行う。
        return par[i]        // 再帰計算で値を返す
    }
    
    mutating func unite(_ i:Int,_ j:Int){
        let (ri,rj) = (root(i),root(j))
        if ri == rj {return} // 根が同じなら何もしない
        par[ri] = rj       // ri(iの根)の根をrjとする
    }
    
    mutating func isSame(_ i:Int,_ j:Int)->Bool { root(i) == root(j) }
}

var uf = UnionFind(10)
print(uf.isSame(0,9)) // false
uf.unite(0,5) // 0と5が同じグループ
print(uf.isSame(0,9)) // false
uf.unite(5,9) // 5と9が同じグループ
print(uf.isSame(0,9)) // true -- 0と9が同じグループ

高機能化

前記のコードで、配列par[i]は、番号iの親ノードの番号を格納していた。
当然、番号iが最上位の親ノード（根）であれば、par[i]には、iが格納される。
だけど、根の要素に根の番号が格納されているのはMOTTAINAI!と、地球上の誰かがふと思った。…地球上の誰かがふと思ったのだ…配列（みんな）の未来を守らねばと…
そんなわけで、iが根であるとき、par[i]を負数とし、その負数の絶対値は根に属するノードの数とすることにします。つまり、根の要素の値を見ることで、そのグループに属する数が分かることになります。
この思想で前記のコードを書き換える。まず、配列parの初期値は全て「-1」となる。つまり、全ての要素が根(負数)であり、その根に属するノード数は根自身のみの1(負数の絶対値)となる。

struct UnionFind {
    
    private var par:[Int] = [] // treeの親ノードを格納する。
    var parent:[Int] { par }
    
    init(_ n:Int) {
        par = [Int](repeating:-1,count:n)
    }
    
    mutating func root(_ i:Int)->Int {
        if par[i] < 0 {return i} // 自身が根なら自身を返す
        par[i] = root(par[i])  // 再帰計算の途中で経路圧縮もついでに行う。
        return par[i]        // 再帰計算で値を返す
    }
    
    mutating func unite(_ i:Int,_ j:Int){
        let (ri,rj) = (root(i),root(j))
        if ri == rj {return} // 根が同じなら何もしない
        
        // riの根をrjとする。
        par[rj] += par[ri] // 根rjの仲間の数に、riの仲間の数を加える
        par[ri] = rj       // ri(iの根)の根をrjとする
    }
    
    mutating func isSame(_ i:Int,_ j:Int)->Bool { root(i) == root(j) }
}

// 入力
var uf = UnionFind(10)
uf.unite(0,5) // 0と5が同じグループ
uf.unite(5,9) // 5と9が同じグループ
uf.unite(2,3) // 2と3が同じグループ
uf.unite(2,4) // 2と4が同じグループ

for i in 0..<10 {
    print(i,uf.parent[i],uf.root(i)) //ノード番号、親番号、根番号
}

// 出力
0 5 9 //ノード0の親は5で、根は9
1 -1 1
2 3 4
3 4 4
4 -3 4
5 9 9 //ノード5の親は9で、根は9
6 -1 6
7 -1 7
8 -1 8
9 -3 9 //ノード9は根で、仲間は3つ

実装したUnion-Findの活用方法

グループ化が高速で実行できたとして、何が嬉しいの？
というのが普通の感想だと思うので、活用例を示そうと思う。
とは言え、活用例もAtCoderの問題を解くことだから、「活用方法」と呼ぶのは恥ずかしいかも。
上記の問題は、最大のお友達集団の人数を調べることが回答に繋がる。最大のお友達集団の全員を別々のグループに分けるための最小のグループ数は、当然、最大のお友達集団の人数と一致する。
入力値

10 4 //ノード数10、お友達ペア4
3 1 //ペア① 3と1は友達
4 1 //ペア② 4と1は友達
5 9 //ペア③ 5と9は友達
2 6 //ペア④ 2と6は友達

回答用コード

let (N,M) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]

var uf = UnionFind(N)

for _ in 0..<M {
    // 問題文のノード番号から1を引いて、配列のindexにする。
    let (a,b) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0]-1,$0[1]-1)}[0] 
    uf.unite(a,b)
}

var ans = 0
for i in 0..<N {
    ans = max(ans,-uf.parent[i])
}

print(ans)

さいごに

活用例で示したAtCoderの問題は、Union-Findを知っていれば、すごく簡単な問題に思えるけど、ABCコンテストが6問だった時代のd問題なんだよね。だから、実は難易度高めの問題。
そんな問題も簡単に解けちゃうUnion-Findは強力なツールじゃないだろうか（やや自信なし）