RubyでUnion-Find木

posts/2019-06-05diary

Rubyで書いたUnion-Findを持っていなかったのでこの記事を書いた。

Kotlinで使っていたものをRubyで書き直しただけなので、
あまりRubyっぽい書き方にはなっていないかもしれない。

そもそもUnion-Find木とはなんぞやということを簡単に書いておく。

Union-Find木 is 何

グループ分けを管理するためのデータ構造で、以下の2つの操作を高速に行える。

2つの要素が同じグループに属しているかを判定する(ある要素がどのグループに属しているか求める)
2つのグループを併合する

正式には、素集合データ構造(disjoint-set data structure)に対する上記の2つの操作をUnion-Findアルゴリズムというらしく、
もしかしたらデータ構造そのものをUnion-Find木と呼ぶのはおかしいのかもしれない。

グループを連結成分と呼んだりもする。

たとえば、

高速道路Aで繋がっている都市1, 都市2, 都市3のグループA
高速道路Bで繋がっている都市4, 都市5, 都市5のグループB
高速道路Cで繋がっている都市6, 都市7, 都市8のグループC

みたいなものがあったときに、

都市5はどのグループに属しているか？
都市1と都市7は同じグループに属しているか？
都市2と都市7が新たに高速道路で結ばれたのでグループを併合する
そうなったとき、都市3と都市4は同じグループに属しているか？

のようなクエリに高速に答えられる。

図で見るととてもわかりやすいので、Union find(素集合データ構造)を参照されたし。

できないこと

グループを分割することはできない

実装

元々使っていたKotlinの実装はこのページを参考にした。
(いつも参考にさせてもらっているけんちょんさんのブログ)
けんちょんの競プロ精進記録 AtCoder ABC 120 D - Decayed Bridges (400 点)

この記事にある通り、併合処理の実装には2種類の方針がある。

merge by rank
merge by size

Union-Find木で検索した時によく見つかるのは「merge by rank」を採用している方で、
Union find(素集合データ構造)のスライドでもrankの方で解説されている。

rankというのはグループを表す木の高さのことで、rankを管理するための専用の配列を持っておく。
parent配列とrank配列の両方を管理する必要があるため少し複雑になるような気がする。

一方sizeでやれば、最初のparent配列の初期化を工夫するだけでよく、実装がシンプルになる。
副産物として、グループの大きさを簡単に取得することもできる。

また親ノードを探すときに、ノードの付け替えを行う経路圧縮というものを行い効率化を図っている。
さらに併合するときにはマージテクと呼ばれるものを使っている。

class UnionFind
    attr_accessor :parent
    
    def initialize(size)
        @parent = Array.new(size, -1)
    end

    # xが所属するグループのルートノードを返す
    def root(x)
        if @parent[x] < 0
            x
        else
            # 経路圧縮を行なう
            @parent[x] = root(@parent[x])
        end
    end

    # xとyが同じグループに属しているか判定する
    def same_parent?(x, y)
        root(x) == root(y)
    end

    # xとyが所属するグループを併合する
    def merge(x, y)
        x_root = root(x)
        y_root = root(y)

        return false if x_root == y_root

        # merge technique
        if @parent[x_root] > @parent[y_root]
            x_root, y_root = y_root, x_root
        end

        @parent[x_root] += @parent[y_root]
        @parent[y_root] = x_root

        return true
    end

    # xが所属するグループのサイズを返す
    def size(x)
        return -@parent[root(x)]
    end
end

動きをおってみる

都市1, 2, 3, 4, 5がある
都市1, 3, 5を繋げてひとつのグループにする
都市2, 4を繋げてひとつのグループにする
この時点で2つのグループがある(奇数都市と偶数都市)
都市4と都市5を繋げる
全ての都市が繋がりひとつのグループになる

uf.parentでグループを管理しており、値によって複数の意味をもつ

値が負の数なら、それはルートノードである
さらに負の数は、その絶対値がグループのサイズを表している
0以上の数は、自分がぶら下がっている親ノードのインデックスを表している

$ irb
irb > require './union_find.rb'
irb > uf = UnionFind.new(5)

irb > uf.parent
=> [-1, -1, -1, -1, -1] # 最初は全てがルート(グループはない)

irb > uf.merge(0, 2) # 都市1と都市3を繋げる
irb > uf.parent
=> [-2, -1, 0, -1, -1] # 都市1がルートになり、都市3は自分がぶら下がっている親ノードのインデックスである0をもつ

irb > uf.merge(2, 4) # 都市3と都市5を繋げる
irb > uf.parent
=> [-3, -1, 0, -1, 0] # 都市5は{都市1, 都市3}のグループに属するようになる

irb > uf.merge(1, 3) # 都市2と都市4を繋げる
irb > uf.parent
=> [-3, -2, 0, 1, 0]

irb > uf.merge(3, 4) # 都市4と都市5を繋げる(偶数都市グループと奇数都市グループを併合)
irb > uf.parent
=> [-5, 0, 0, 1, 0] # マージテクにより、小さい方のグループが大きい方のグループのルートに繋がれる

グループのサイズを知りたい場合、例えば都市4が属するグループのサイズを知りたいときは

都市4がぶら下がっている親ノードは都市2であることがわかる(1はインデックスなので)
都市2がぶら下がっている親ノードは都市1である(0はインデックスなので)
都市1の値は負の数なのでこれがルートで、その絶対値である5がグループのサイズ
よって都市4が属するグループのサイズは5であることがわかる

少しわかりづらいかもしれないけどこんな感じ。