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]

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

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

次に、要素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木 を使って問題を解いてみます。