[Ruby] UnionFind木の実装

Union-Findアルゴリズムを木構造で実装する。
wikiに擬似コードでアルゴリズムが書いてある。

Union-Findがどんなもんかというと、何個かの要素があるとして

2番目と5番目を同じグループにして(Union
3番目と4番目を同じグループにして(Union
2番目と3番目を同じグループにした(Union
さあ、2番めと4番目は同じグループだろうか?(Find

という操作を、非常に高速に行える。
かっこ良く言うと、計算量はアッカーマン関数の逆関数で高速らしい。なにそれかっこいい!

まずRubyのコードを載せる。

class Node
  attr_accessor :parent, :rank

  def initialize(n)
    @parent = n
    @rank = 0
  end
end

class UnionFindTree
  def initialize(n)
    @nodes = (0..n).to_a.map { |i| Node.new(i) }
  end

  def find(x)
    return x if @nodes[x].parent == x

    return @nodes[x].parent = find(@nodes[x].parent)
  end

  def unite(a, b)
    a = find(a)
    b = find(b)
    return if a == b

    if @nodes[a].rank < @nodes[b].rank
      @nodes[a].parent = b
    else
      @nodes[b].parent = a
      @nodes[a].rank += 1 if @nodes[a].rank == @nodes[b].rank
    end
  end

  def same?(a, b)
    find(a) == find(b)
  end

  # 確認用。アルゴリズムとは関係無い
  def parents
    @nodes.map(&:parent)
  end
end

使い方は簡単。
まずUnionFindTree.newで集合を作成する。

tree = UnionFindTree.new(10)

10個ノードが作成される。
これらは初期状態では独立しているので、parentは自分を指している状態になる。
それぞれのparentだけを見てみるとこんな感じ

tree.parents
# => [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

Union操作をしてみる。この実装ではuniteメソッド

tree.unite(2, 5)

tree.parents
# => [0, 1, 2, 3, 4, 2, 6, 7, 8, 9, 10]

2番目と5番目をグループ化したので、5番目のparentが2になっている。
もっとuniteしてみる

tree.unite(3,4)
tree.parents
# => [0, 1, 2, 3, 3, 2, 6, 7, 8, 9, 10]

tree.unite(2,3)
tree.parents
# => [0, 1, 2, 2, 3, 2, 6, 7, 8, 9, 10]

まず3番目と4番目をグループ化した。4番目のparentは3番目になる。
その後2番めと3番目をグループ化したので、3番目のparnetが2番目になっている。問題なさそう。

Union-Find木を高速化する方法として、Find時に直接親に紐づけてしまうというものがある。
この実装が正しければ、tree.find(4)を実行したら4番目のparentはもともと3番目だったものが2番目を指すようになるはず。。

実行してみる

tree.find(4)
tree.parents
# => [0, 1, 2, 2, 2, 2, 6, 7, 8, 9, 10]

4番目のparentは2番目になった。
うん、良いね。

2015年3月9日追記

「uniteメソッドが間違っているのでは」という指摘をいただきました!たしかにおかしかったので修正しました。ありがとうございます!

言い訳をすると、ブログにRubyでアルゴリズムを書くときには、C++で実装してからRubyに書き直すという手順を踏むんだけど(学習のため2回別言語でアルゴリズムを書くようにしている)、C++版では間違ってなかったのに、Ruby版を作った時に失敗したようだ。。
もっとテストケースを増やしてやらなければ。。

タイトルとURLをコピーしました