Union-Find木

問題

頂点数 $N$ のグラフに対して、次の２種類のクエリーを処理してください。

　クエリー１：頂点 $u$ と頂点 $v$ を双方向に結ぶ辺を追加する。
　クエリー２：頂点 $u$ と頂点 $v$ は同じ連結成分に属するかを答える。

最初はグラフに辺が一本もなく、与えられるクエリ―の数は $Q$ 個であるとします。

[制約]
🔹$ 2 \leqq N \leqq 100000 $
🔹$ 1 \leqq Q \leqq 100000 $

解き方・ソースコード

この問題は、Union-Find木というデータ構造を使って解くことができます。

Union-Fin木dとは、グループ分けを効率的に管理することができるデータ構造です。

次のような２種類のクエリ―を高速に処理することができます。

🔹統合クエリー：要素 $u$ を含むグループと、要素 $v$ を含むグループを統合する。
🔹回答クエリー：要素 $u$ と 要素 $v$ が同じグループにあるかどうかを判定する。

以下のソースコードは、Union-Find木を使って与えられたグラフの連結成分を管理し、クエリーを処理します。

入力例１として、3つの頂点と4つのクエリーを設定しています。

クエリー１は頂点1と頂点2を結ぶ辺を追加し、クエリー２は頂点2と頂点3が同じ連結成分に属するかを判定するものです。

まず、Union-Find木のクラスを定義します。

コンストラクタでは、頂点数 n を受け取り、par と size の初期化を行っています。

par は親の頂点を保持する配列で、初期値は全て -1 です。

size はグループの頂点数を保持する配列で、初期値は全て 1 です。

次に、root関数では、引数で指定された頂点 x の根を探索して返します。

par[x] が -1 になるまで、頂点 x の親に移動していきます。

unite関数は、引数で指定された要素 u, v を同一のグループに統合します。

まず、それぞれの根を取得します。

もし根が異なる場合は、グループの頂点数が小さい方を大きい方に統合します。

par 配列には親の頂点を、size 配列にはグループの頂点数を記録しています。

same 関数は、引数で指定された要素 u, v が同一のグループに属するかどうかを判定します。

関数内ではroot 関数を呼び出して u, v の根を比較します。

最後に、クエリーの処理を行います。

クエリー１の場合は、union 関数を呼び出して頂点 u, v を統合します。

クエリー２の場合は、same 関数を呼び出して頂点 u, v が同一のグループに属するかどうかを判定し、結果に応じてメッセージを出力します。

[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
50
51
52
53
#-------- 入力例１ ---------
N = 3       # 頂点数
Q = 4       # クエリ―数
queries = [
    (1, 1, 2),
    (2, 1, 3),
    (1, 2, 3),
    (2, 2, 3)
]
#---------------------------
# Union-Find 木
class unionfind:
    # n 頂点の Union-Find 木を作成
    # 頂点番号が 1 - indexed になるように実装するので n + 1 としている。
    def __init__(self, n):
        self.n = n
        self.par = [ -1 ] * (n + 1) # 最初は親が無い
        self.size = [ 1 ] * (n + 1) # 最初はグループの頂点数が 1

    # 頂点 x の根を返す関数
    def root(self, x):
        # 1 個先（親）がなくなるまで、1 個先（親）に進み続ける
        while self.par[x] != -1:
            x = self.par[x]
        return x

    # 要素 u, v を統合する関数
    def unite(self, u, v):
        rootu = self.root(u)
        rootv = self.root(v)
        if rootu != rootv:
            # u と v が異なるグループのときに処理を行う
            if self.size[rootu] < self.size[rootv]:
                self.par[rootu] = rootv
                self.size[rootv] += self.size[rootu]
            else:
                self.par[rootv] = rootu
                self.size[rootu] += self.size[rootv]

    #  要素 u と v が同一のグループかどうかを返す関数
    def same(self, u, v):
        return self.root(u) == self.root(v)

# クエリの処理
uf = unionfind(N)
for tp, u, v in queries:
    if tp == 1:         # クエリ―１の処理を行う
        uf.unite(u, v)
    if tp == 2:         # クエリ―２の処理を行う
        if uf.same(u, v):
            print(f'頂点 {u} と頂点 {v} は同じ連結成分に属する。')
        else:
            print(f'頂点 {u} と頂点 {v} は同じ連結成分に属さない。')

[実行結果（入力例１）]

頂点 1 と頂点 3 は同じ連結成分に属しない。

頂点 2 と頂点 3 は同じ連結成分に属する。

クエリ―１で頂点またはそのグループを結合し、クエリ―２の回答として各頂点が同じグループに属するかどうかを判定することができました。