けんちょん本（「アルゴリズムとデータ構造」）の第11章で扱われているUnion-Findについて、Pythonで実装してみました。
前提知識：素集合データ構造に対する操作
素集合データ構造とは、互いに重なり合うことが無い集合を持つデータ構造（おおざっぱにいうときれいに分類されているということ）。このデータ構造に対して、Union（2つの集合を1つに統合する）とFind（特定の要素がどの集合に属しているかを調べる）という操作を行うアルゴリズムをUnion-Findアルゴリズムと呼びます。この2操作に、MakeSet（集合を作成する）を加えることにより、さまざまな分割問題をとけるようなツールとして利用できます。
Pythonで実装してみた
けんちょん本とはちょっと実装方針を変えています。
- 素集合データ構造（DisjointSet）の中にUnion-Findを含めた。
- Root→Find、Unite→Unionにメソッド名を変えた。
class DisjointSet:
def __init__(self) -> None:
self.parent = {}
self.size = {}
def makeSet(self, set):
for e in set:
self.parent[e] = e
self.size[e] = 1
def Find(self, x):
if self.parent[x] == x:
return x
else:
# 経路圧縮
self.parent[x] = self.Find(self.parent[x])
return self.parent[x]
def Union(self, x, y):
a = self.Find(x)
b = self.Find(y)
if a == b:
return False
# union by sizeで効率化
if self.size[a] < self.size[b]:
a, b = b, a
self.parent[b] = a
self.size[a] += self.size[b]
return True
def issame(self, x, y):
return self.Find(x) == self.Find(y)
def getsize(self, x):
return self.size[self.Find(x)]
def printSet(set, ds):
return [ds.Find(i) for i in set]
def test_unionfind():
ds = DisjointSet()
ds.makeSet([1, 2, 3, 4, 5, 6])
ds.Union(1, 2)
ds.Union(2, 3)
ds.Union(5, 6)
assert ds.issame(1, 3) == True
assert ds.issame(2, 5) == False
ds.Union(1, 6)
assert ds.issame(2, 5) == True