ruby で union-find

rubyでunion-findを実装する

最近やっとグラフについて入門し始めたのですが、その第一弾としてunion-findが出てきました。
特徴としては、

• 要素同士をグループ分けできる
• 同じグループに所属するか判定できる

というように、それぞれの要素をグループで分けて見ることができるというものです。
早速実装を見ていきます。

union-find

# Union-Find
class UnionFind
  attr_reader :par, :siz

  # 初期化
  def initialize(n)
    # 各要素の親ノードを保持する
    @par = Array.new(n, -1)
    # グループのサイズを保持する
    @siz = Array.new(n, 1)
  end

  # 根を求める
  def root(x)
    if par[x] == -1
      return x
    else
      # 経路圧縮のため、rootを再起的に呼び出し、-1になるまで行う
      return par[x] = root(par[x])
    end
  end

  # x, y が同じグループに属するかどうか（根が一致するかどうか）
  def same?(x, y)
    return root(x) == root(y)
  end

  # x を含むグループと y を含むグループとを併合
  def unite(x, y)
    # それぞれの根を取得
    x = root(x)
    y = root(y)

    # 同じグループであれば早期リターン
    return false if x == y

    # y 側のサイズが小さくなるようにする
    x, y = y, x if siz[x] < siz[y]

    # y を x　の子とする
    par[y] = x
    siz[x] += siz[y]
    return true
  end

  # グループのサイズを返す
  def size(x)
    return siz[root(x)]
  end
end

細かく見ていきます。

attr_reader :par, :siz

よく見るインスタンス変数およびアクセサの設定です。
簡単にいうと、インスタンス変数に外部からアクセスできるメソッドが自動で利用できるようになります。（javaでいう getter, setter と同じかそれと同等の働きをするものです。）
par は 各要素の親を保持し、 siz はそれぞれの要素が所属するグループの要素数（サイズ）を保持します。

  # 初期化
  def initialize(n)
    # 各要素の親ノードを保持する
    @par = Array.new(n, -1)
    # グループのサイズを保持する
    @siz = Array.new(n, 1)
  end

初期化メソッドです。nには要素の数を入れます。
parの初期値の-1は、親を表します。最初はどの要素も単独で存在しているので、全てに-1を入れます。
sizも同様、最初はどの要素も単独で存在しているので、1を入れます。

  # 根を求める
  def root(x)
    if par[x] == -1
      return x
    else
      # 経路圧縮のため、rootを再起的に呼び出し、-1になるまで行う
      return par[x] = root(par[x])
    end
  end

root関数の実装です。要素の根（最終的な親）を求める関数になります。
-1の場合、親ということなのでそのまま自身の数を返します。
そうでない場合、root関数を再起的に呼び出し、-1となるまで繰り返します。

  def same?(x, y)
    return root(x) == root(y)
  end

same?関数の実装です。関数名はrubyっぽくしています。
先ほど定義したroot関数を呼び出し、親要素が同じであれば同じグループということになるため、それぞれの根を比較しています。

  # x を含むグループと y を含むグループとを併合
  def unite(x, y)
    # それぞれの根を取得
    x = root(x)
    y = root(y)

    # 同じグループであれば早期リターン
    return false if x == y

    # y 側のサイズが小さくなるようにする
    x, y = y, x if siz[x] < siz[y]

    # y を x　の子とする
    par[y] = x
    siz[x] += siz[y]
    return true
  end

続いて、unite関数の実装です。
最初に与えられた要素のrootを求めます。同じであればすでに同じグループに属しているので、特に何もせず早期リターンします。
そうでない場合、siz関数を利用し要素のサイズを調整します。（x,yについてxのグループが大きくyのグループが小さいという状態にします。）
そして、小さい方（y）の親をxとし（yをxに併合し）、グループのサイズを変更します。
詳細は省きますが、この小さいグループを大きいグループに併合する、というのは 「union by size」と呼ばれるもので、計算量を減らすのに一役買っています。

最後に、実際の動きを見てみます。

# テスト
uf = UnionFind.new(7) # {0}, {1}, {2}, {3}, {4}, {5}, {6}
p uf.par # [-1, -1, -1, -1, -1, -1, -1]

uf.unite(1, 2) # {0}, {1, 2}, {3}, {4}, {5}, {6}
uf.unite(2, 3) # {0}, {1, 2, 3}, {4}, {5}, {6}
uf.unite(5, 6) # {0}, {1, 2, 3}, {4}, {5, 6}

p uf.par         # [-1, -1, 1, 1, -1, -1, 5]
p uf.same?(1, 3) # true
p uf.same?(2, 5) # false

uf.unite(0, 1) # {0, 1, 2, 3}, {4}, {5, 6}
# ※union-by-sizeにより、二つ目の要素が１つ目の要素の親となっている
p uf.par       # [1, -1, 1, 1, -1, -1, 5]
p uf.size(0)   # 4
p uf.size(1)   # 4
p uf.size(2)   # 4
p uf.size(3)   # 4
p uf.size(4)   # 1
p uf.size(5)   # 2
p uf.size(6)   # 2

以上で終了です。
割と内部での定義や処理は初歩的なメソッドで実装できる、というのが面白いです。
グラフというと急に身構えてしまうのですが、union-findは利用・実装も簡単なので入りやすい気がしました。