1. はじめに
競技プログラミングをしている方なら、一度はUnion-Find木なるものを聞いたことがあると思います。
今回は、AtCoderで過去問を解いている時に、Union-Find木を使って解くと分かりやすい問題があったので、備忘録程度にUnion-Find木の使い方を書いていきたいと思います。
この記事はAtCoderのレートで茶色〜緑色前半くらいの方向けに書いているつもりです。
2. Union-Find木とは
class UnionFind():
def __init__(self, n):
self.n = n
self.parents = [-1] * n
def find(self, x):
if self.parents[x] < 0:
return x
else:
self.parents[x] = self.find(self.parents[x])
return self.parents[x]
def union(self, x, y):
x = self.find(x)
y = self.find(y)
if x == y:
return
if self.parents[x] > self.parents[y]:
x, y = y, x
self.parents[x] += self.parents[y]
self.parents[y] = x
def size(self, x):
return -self.parents[self.find(x)]
def same(self, x, y):
return self.find(x) == self.find(y)
def members(self, x):
root = self.find(x)
return [i for i in range(self.n) if self.find(i) == root]
def roots(self):
return [i for i, x in enumerate(self.parents) if x < 0]
def group_count(self):
return len(self.roots())
def all_group_members(self):
return {r: self.members(r) for r in self.roots()}
def __str__(self):
return '\n'.join('{}: {}'.format(r, self.members(r)) for r in self.roots())
上記のコードはこちらから引用しました。
Union-Find木について書かれているページを見ると、Union-Find木は「グループ分けを木構造で管理するデータ構造のこと」と書かれています。
Union-Find木の主な処理としては、
・union(a,b)
・find(a)
の二つで、Unionでは要素aが所属するグループと要素bが所属するグループを併合させる処理であり、
Findでは要素aが所属するグループの根を返す処理をします。
Findの処理は経路圧縮という、同じグループに所属しているのかを判別する時に高速化する処理に用いられます。
経路圧縮とは、下の図のように深く伸びている木構造を、
各要素から根まで１ステップでアクセスできるように変える処理のことです。
この２つの処理を組み合わせて以下のようなコードを書きました。
uf = UnionFind(6)
print(uf.parents)
# [-1, -1, -1, -1, -1, -1]
uf.union(0, 1)
print(uf.parents)
# [-2, 0, -1, -1, -1, -1]
uf.union(1, 2)
print(uf.parents)
# [-3, 0, 0, -1, -1, -1]
uf.union(2, 3)
print(uf.parents)
# [-4, 0, 0, 0, -1, -1]
uf.union(4, 5)
print(uf.parents)
# [-4, 0, 0, 0, -2, 4]
uf.union(0, 4)
print(uf.parents)
# [-6, 0, 0, 0, 0, 4]
uf.find(5)
print(uf.parents)
# [-6, 0, 0, 0, 0, 0]
上の例では、大きさ6のUnion-Find木を作成しています。
先ほど説明したように、unionメソッドで要素xと要素yを併合させています。
途中で要素5の親要素が4から0に変わっています。
これは、findメソッドにより経路圧縮を行うことで、要素5の親要素である4の親要素0に
親要素が変化したことを表しています。
3. 解いた問題
今回はこちらの問題を解きました。
AtCoder Begginer Contest 097 D. Equals
こちらの問題は、swapできる整数のペアを木構造で考えていくと、
同じ木構造に属している整数同士は自由に行き来することができると考えられます。
Union-Find木のデータ構造を用いれば、同じグループ(木構造)に属しているかどうかを高速で判断することができるため、この問題を解くことができます。
4. ソースコード
class UnionFind():
def __init__(self, n):
self.parent = [-1 for _ in range(n)]
# 正==子: 根の頂点番号 / 負==根: 連結頂点数
def find(self, x):
if self.parent[x] < 0:
return x
else:
self.parent[x] = self.find(self.parent[x])
return self.parent[x]
def unite(self, x, y):
x, y = self.find(x), self.find(y)
if x == y:
return False
else:
if self.size(x) < self.size(y):
x, y = y, x
self.parent[x] += self.parent[y]
self.parent[y] = x
def same(self, x, y):
return self.find(x) == self.find(y)
def size(self, x):
x = self.find(x)
return -self.parent[x]
def is_root(self, x):
return self.parent[x] < 0
n, m = map(int, input().split())
p = list(map(int, input().split()))
count = 0
pair_list = []
for i in range(m):
array = list(map(int,input().split()))
array[0] -=1 ; array[1] -= 1
pair_list.append(array)
uf = UnionFind(len(p))
# 整数のペアを同じグループに結合する
for i in range(m):
uf.unite(pair_list[i][0], pair_list[i][1])
# 経路圧縮をして、親要素を根に更新する
for i in range(n):
uf.find(i)
# 順列pに含まれる要素piの値とindex値(+1)が一致、もしくは
# index[i] とindex[p[i]] が同じグループに所属しているときにcount+=1
for i in range(n):
if p[i] == i+1:
count += 1
else:
if uf.same(p[i]-1, p[p[i]-1]-1) is True:
count += 1
else:
continue
print(count)
前半のUnion-Findのクラスは誰かの提出コードから拝借しました。
参考文献
・Union-Find木の解説と例題
・Union-Find(素集合データ構造)
・PythonでのUnion-Find（素集合データ構造）の実装と使い方