はじめに
Union-Findが初耳だという方は
https://www.slideshare.net/iwiwi/ss-3578491
https://www.slideshare.net/chokudai/union-find-49066733
こちらのスライドを御覧ください。
本記事では簡単な実装と、
実際にAtCoderの問題を解く時に追加実装する機能について記します
Union-Findとは？
A,B,C,...,Xさんがいます。
「AさんとBさんは友達！」
「CさんとEさんは友達！」
...
という交友関係の情報が与えられます。
友達も、友達の友達も、友達の友達の友達も、... すべて友達だとみなす時、
「○さんと○さんは友達ですか？」
という質問に高速に答えられるようなデータ構造です。
これ以上詳しく知りたい方は、上記スライドぜひご覧ください。
(図などもあり超わかりやすいです)
実装
突然ですが実装に移ります。
class UnionFind
def initialize(size)
@rank = Array.new(size, 0)
@parent = Array.new(size, &:itself)
end
def unite(id_x, id_y)
x_parent = get_parent(id_x)
y_parent = get_parent(id_y)
return if x_parent == y_parent
if @rank[x_parent] > @rank[y_parent]
@parent[y_parent] = x_parent
else
@parent[x_parent] = y_parent
@rank[y_parent] += 1 if @rank[x_parent] == @rank[y_parent]
end
end
def get_parent(id_x)
@parent[id_x] == id_x ? id_x : (@parent[id_x] = get_parent(@parent[id_x]))
end
def same_parent?(id_x, id_y)
get_parent(id_x) == get_parent(id_y)
end
end
最低限4つのメソッドを実装します。
1.initialize初期化
2.unite ノードxとノードyを連結させる(=同じグループにする)
3.get_parent 親ノードの番号を取得する
4.same_parent? ノードx, yが連結か確かめる
initialize
UnionFind.new としてインスタンスを作ったときにそのインスタンスに対して呼び出されるメソッドです。
UnionFind のインスタンスには、各ノードのrankと、各ノードの親ノードの番号を持っていてほしいので、上記実装になっております。
Array.new(size) {|idx| block }引数objの代わりにブロックを渡すと、sizeの数だけブロックを繰り返し、
ブロックが返す値を配列の要素にします。ブロック引数idxには要素の位置が整数で入ります。
連番で初期化したいときは上記仕様を利用すると楽ちんですね。
追記
コメントでいただきましたArray.new(size, &:itself) のほうがかっこいいと感じたのでそれを採用させていただきました。
[参考]
https://ref.xaio.jp/ruby/classes/array/new
https://ref.xaio.jp/ruby/classes/object/initialize
unite
なぜこういう処理をするかは、最初にリンクを張りましたスライドを御覧いただきたいです。
rankが大きいものに連結していくのは、そうしたほうが効率が良いからです。
get_parent
一番上のノードの番号(=親ノードの番号)を取得します。
このメソッドが呼び出されたときに、
ちゃんとひ孫から始まったらひ孫自身も孫も子も@parentの値を更新するために、
上記実装になっています。
same_parent?
get_parentを使って、2つのノードの親が同じか調べます
実際に問題を解く1
https://atc001.contest.atcoder.jp/tasks/unionfind_a
上記クラスを使えば、あとは入力を受け取って適当に出力するだけです。
# ノードの数、クエリの数を受け取る
nn, qq = gets.split.map(&:to_i)
union = UnionFind.new nn
# クエリの回数だけ、入力受け取りを回す
qq.times do
k = gets.split.map(&:to_i)
if k[0] == 0 # 0なら結合操作を行う
union.unite(k[1],k[2])
else # 1なら連結か判定して出力をする
puts union.same_parent?(k[1],k[2]) ? 'Yes' : 'No'
end
end
AC => https://atc001.contest.atcoder.jp/submissions/5677828
https://atc001.contest.atcoder.jp/submissions/5682673
実際に問題を解く2
https://arc032.contest.atcoder.jp/tasks/arc032_2
N個ノードがあって結合クエリがM件入力される。
最終的にいくつのグループに別れていますか？
1.今まで通り結合クエリを処理する(前述と同じ)
2.いくつのグループに分かれているか調べる
2をやるためには、@parentのすべての親を確かめて、
それが何種類あるか確かめる必要があります。
@parent それぞれにget_parentを適用した結果を、
uniq(重複要素を取り除く)して、そのsizeを見れば大丈夫です。
class UnionFind
# 略
def size
@parent.map { |id_x| get_parent(id_x) }.uniq.size
# @parent.map.with_index.count(&:==)
# とすることも出来ます!
end
end
size を何度も問われる可能性があるときには
@parent.map { |id_x| get_parent(id_x) }の結果を
インスタンス変数に格納してもいいと思います。
(次からは @size.map { |id_x| get_parent(id_x) } とできるはずです。)
nn, mm = gets.chomp.split.map(&:to_i)
union = UnionFind.new nn
mm.times do
union.unite(*gets.chomp.split.map { |g| g.to_i - 1 })
end
p union.size - 1
AC => https://arc032.contest.atcoder.jp/submissions/5682769
`Setを使う場合(旧)`
こういう集合を調べたいときはSetを使いましょう
https://docs.ruby-lang.org/ja/latest/class/Set.html
class UnionFind
require 'set'
# 略
def size
Set.new(@parent.map { |id_x| get_parent(id_x) }).size
end
end
この問題は、最後にsizeを求めるだけなので不必要ですが、
クエリの途中に何度もsizeを求める必要がある場合は、
このSetをインスタンス変数に持っておいたほうがいいです。
nn, mm = gets.split.map(&:to_i)
union = UnionFind.new nn
1.step(mm, 1) do |_i|
union.unite(*gets.split.map { |g| g.to_i - 1 })
end
p union.size - 1
さいごに
ということでいつUnion-FindがAtCoderで出ても安心ですね。
せっかく作ったのではやくABCに挑戦したいです。
拙い説明にご付き合いいただきありがとうございました。
@scivola さんコメントでいろいろ教えていただきありがとうございました。