Union Find木

Union Find木

Union Find木 は、グループ分けを管理 することができるデータ構造です。

次のようなことを効率的に行うことができます。

  • 要素aと要素bが 同じグループ に属しているかどうかをチェックする。
  • 要素aと要素bの グループを一つ にまとめる。

(グループ同士をまとめることはできますが、要素をグループから除いたり、グループを分けたりすることはできません)

Union Find木 を実現するためのクラスは以下のようになります。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
from collections import defaultdict

# n個の要素を0 ~ n - 1のインデックスで管理する
class UnionFind():
def __init__(self, n):
self.n = n # 要素の数
self.parents = [-1] * n # 各要素の親要素の番号を格納するリスト

def find(self, x): # 要素[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]が属するグループと要素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): # 要素[x]が属するグループのサイズ(要素数)を返す
return -self.parents[self.find(x)]

def same(self, x, y): # 要素[x],要素[y]が同じグループに属するかどうかを返す
return self.find(x) == self.find(y)

def members(self, x): # 要素[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): # 全グループの全要素を返す
group_members = defaultdict(list)
for member in range(self.n):
group_members[self.find(member)].append(member)
return group_members

def __str__(self): # printでの表示用。ルート要素: [そのグループに含まれる要素のリスト]を文字列で返す
return '\n'.join(f'{r}: {m}' for r, m in self.all_group_members().items())

動作確認

実装した Union Find木 の動作を確認してみます。

要素数5個(0,1,2,3,4)のUnion Find木を定義し、[0,1] のグループと [2,3,4] のグループに分けます。

[Google Colaboratory]

1
2
3
4
5
6
uf = UnionFind(5)
uf.union(0, 1)
uf.union(2, 3)
uf.union(3, 4)

print(uf)

[実行結果]

0: [0, 1]

2: [2, 3, 4]

2つのグループに分けることができました。

親の要素(一番もとになる要素) がグループの キー となっていることが確認できます。

次に、要素1と要素3が所属するグループを調べてみます。

[Google Colaboratory]

1
2
print(uf.find(1))
print(uf.find(3))

[実行結果]

0

2

要素1がグループ0に所属し、要素3がグループ2に所属していることが分かりました。

今度は、要素0,要素2が同じグループに属するかどうかと、要素3,要素4が同じグループに属するかどうかを調べます。

[Google Colaboratory]

1
2
print(uf.same(0, 2))
print(uf.same(3, 4))

[実行結果]

False

True

要素0,要素2は違うグループであり、要素3,要素4は同じグループに属することを確認できました。

最後に、要素3が所属するグループの全要素を表示します。

[Google Colaboratory]

1
print(uf.members(3))

[実行結果]

[2, 3, 4]

要素3が所属するグループの全要素は、2,3,4であることが分かりました。

次回は、Union Find木 を使って問題を解いてみます。