UnionFindの改善を今更やった

UnionFindの定番らしい計算量改善(経路圧縮、union by size)をやってみた。

3ヶ月くらい前にUnionFindを実装したのだけど、その時はUnionFindの高速化対応をサボっていた。当時取り組んでいた問題はACしたし、「応用はちょっと難しそうだな」とか謎に尻込みしていた気がする。

今回、クラスカル法の精進でAOJの最小全域木問題に先述のUnionFind使ったところ、あえなくTLEしてしまったので改善することにした。

UnionFindそのものについては過去の記事に書いていて、なんなら改善方法含めてWeb上にわかりやすい情報が山程あるのであえて書かないでおく(書き始めると時間が溶けるのでサボりたいのが本音)。

改善方法

UnionFindはノード同士が同じグループに属しているかを管理するデータ構造で、「あるノードの根は何か？」を再帰的に調べる処理が重い。普通にノードの結合を続けていくと、N個のノードで根を調べる際の最悪計算量がN-1になってしまう(要素数Nで木の高さがNになった場合)。

なので木の高さを低く保つことができれば計算量が改善できるということ。その方法として、経路圧縮とunion by size(rank)がある。どちらか一方でも行えばO(logN)になり、両方組み合わせるとアッカーマン関数の逆関数αを用いてO(α(N))で計算できる(正直よくわかってないけどO(logN)よりも早い)らしい。

一応改善前のUnionFindを置いておく。

class UnionFind
  def initialize(n)
    # par[i]: iの親要素番号を表す
    # par[i] = i の時、par[i]は根
    @par = []
    n.times { |i| @par[i] = i }
  end

  # データxが属する木の根を返す
  def root(x)
    # x自身が根の場合
    return x if @par[x] == x
    # 自身が根でない場合は再起で親を辿る
    return root(@par[x])
  end

  # xとyの木を併合する
  def unite(x, y)
    # xとyの根を取得する
    rx = root(x)
    ry = root(y)
    return if rx == ry # もともと同じ木の場合は何もしない
    @par[rx] = ry # xの親をyにする = yが根になる
  end

  # xとyが同じ木に属しているかを調べる
  def same?(x, y)
    rx = root(x)
    ry = root(y)
    return rx == ry
  end
end

経路圧縮

UnionFindはノード同士が同じ家系図にいるかどうかさえわかればよく、何親等かとか兄弟構成については問われない。なので、Aさんの祖先(家系図の一番上)を調べたタイミングでAさんを祖先の実子であることにしてしまえば良い。この時点での計算量は変わらないが、後からAさんの子であるBさんの祖先を探す時は2つ辿るだけで根にたどり着ける。

実装は驚くほどに簡単。再起で根を見つけた後、根を自分の親に設定するだけ。

  # データxが属する木の根を返す
  def root(x)
    # x自身が根の場合
    return x if @par[x] == x
    # 自身が根でない場合は再起で親を辿る
-   return root(@par[x])
+   return @par[x] = root(@par[x])
  end

これでどれくらい計算量が改善されるか、AOJのテストケースに対してroot(x)の呼び出し回数をカウントしてみた。

❯ ruby GRL_2_A.rb < in.txt
215958829

❯ ruby GRL_2_A.rb < in.txt
541182

3桁ぐらい変わってる。

union by size

経路圧縮がroot(x)を改善するアプローチだったが、union by size(rank)ではノード結合のunite(x,y)に手を入れる。

ある木と木を結合する際に木のサイズを比較し、小さい方の根を大きい方の根の子として結合するようにする(テキストだととてもわかりにくい)。こうすることで、結合後の木の高さの増加分を0または1に抑えることができる。

これも実装は簡単で、グループごとのサイズを持つ変数を追加し、サイズが小さい方の木が大きい方に属するように変えるだけ。(Rubyは2変数の入れ替えをtmpなしでできるからいいな、と思うなどした)

ちなみに木のサイズではなく、木の高さで判断するのをunion by rankというらしい。

  def initialize(n)
    # par[i]: iの親要素番号を表す
    # par[i] = i の時、par[i]は根
    @par = []
+   @siz = Array.new(n,1)
    n.times { |i| @par[i] = i }
  end

  # データxが属する木の根を返す
  def root(x)
    # x自身が根の場合
    return x if @par[x] == x
    # 自身が根でない場合は再起で親を辿る
    return root(@par[x])
  end

  # xとyの木を併合する
  def unite(x, y)
    # xとyの根を取得する
    rx = root(x)
    ry = root(y)
    return if rx == ry # もともと同じ木の場合は何もしない
+   rx,ry = ry,rx if @siz[rx] > @siz[ry] # 必ず大きい方が根になるようにする
    @par[rx] = ry # xの親をyにする = yが根になる
+   @siz[ry] += @siz[rx]
  end

union by sizeだけで改善した後の計算量。

❯ ruby GRL_2_A.rb < in.txt
215958829

❯ ruby GRL_2_A.rb < in.txt
612808

これも3桁改善されてる。

ちなみに経路圧縮とunion by sizeの両方改善した場合。

❯ ruby GRL_2_A.rb < in.txt
422981

桁は変わらないものの改善されてる。

UnionFind自体も実装がとてもシンプルで、データ構造を学ぶ面白さが感じられるなーと思っていたけど、ちょっとした工夫でこれほど計算量が改善されるというのもまた興味深い。

今日は眠いのでそろそろ寝る。