Union-Find木②

問題

ある地域には $N$ 個の駅と $M$ 本の鉄道路線があります。

駅には $1$ から $N$ までの番号が付けられていて、$i$ 本目の路線は駅 $A_i$ と駅 $B_i$ を双方向に結んでいます。

この地域に台風が上陸すると、いくつかの路線は運休になってしまいます。

以下のようなクエリ―が $Q$ 個与えられるのでそのクエリ―を処理してください。

　クエリー１：$x$ 本目の路線が運休になる。
　クエリー２：運休状況を踏まえて駅 $s$ から駅 $t$ へ移動できるかどうかを答える。

解き方・ソースコード

この問題も前回記事同様、Union-Find木を使って解いていきますが、辺を減らしながらの連結状況を答えるのは難しいので、クエリーを逆向きにして辺を増やしながら処理することにします。

まずソースコードの冒頭では、次のように入力データを設定しています。

　・N： 頂点の数（駅の数）。
　・M： 辺の数（２つの駅を結ぶ路線の数）。
　・edges： 辺のリスト。各要素は、頂点の番号のペアで表される。
　・Q： クエリ―の数。
　・query： クエリ―のリスト。各要素は、1または2の数字と、いくつかの頂点の番号の組み合わせで構成される。

このプログラムは、以下のような処理を行います。

　１．辺のリストをインデックス0から始まるように修正する。
　２．クエリ―に応じて、最後まで残っている辺のリストを取得する。
　３．Union-Find木を用いて、最後まで残っている辺でグラフを構成する。
　４．クエリ―を逆順に処理し、連結を増やしながら移動可能かどうかの回答を求める。
　５．回答を出力する。

Union-Find木を使ってグラフを構成していきます。

Union-Find木は、グループを管理するためのデータ構造です。

各要素は、自身が属するグループの根の番号を覚えています。

グループをマージするときには、根の番号を統一します。

これによって、任意の2つの要素が同じグループに属するかどうかを簡単に判定することができます。

このプログラムは、最後まで残っている辺でグラフを構成します。

そして、クエリ―を逆順に処理しながら、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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
#-------- 入力例１ ---------
N = 5       # 駅の数
M = 5       # ２つの駅をつなぐ路線の数
edges = [
    (1, 2),
    (2, 3),
    (3, 4),
    (4, 5),
    (3, 5)
]

Q = 5       # クエリ―数
query = [
    [2, 1, 5],
    [1, 3],
    [2, 1, 5],
    [1, 5],
    [2, 1, 5]
]
#---------------------------
edge = []
# インデックスを0スタートにする
for A, B in edges:
    edge.append((A - 1, B - 1))

# 最後まで残っている辺を求める
last = [True] * M
for q in query:
    if q[0] == 1:
        q[1] -= 1
        last[q[1]] = False
    else:
        q[1] -= 1
        q[2] -= 1

# union-find木
uf = [-1] * N
def root(i):
    while True:
        if uf[i] < 0:
            return i
        if uf[uf[i]] < 0:
            return uf[i]
        uf[i] = uf[uf[i]]
        i = uf[i]

def unite(i, j):
    i = root(i)
    j = root(j)
    if i == j:
        return
    # union by size
    if uf[i] > uf[j]:
        i, j = j, i
    uf[i] += uf[j]
    uf[j] = i

ans = []
for i in range(M):
    if last[i]:
        A, B = edge[i]
        unite(A, B)

# クエリ―を逆順に処理する
for q in reversed(query):
    if q[0] == 1:
        A, B = edge[q[1]]
        unite(A, B)
    else:
        _, U, V = q
        ans.append("移動可能" if root(U) == root(V) else "移動不可")

# クエリ―の順番を戻して、回答を表示
for s in reversed(ans):
    print(s)

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

移動可能

移動可能

移動不可

クエリ―に応じて、駅間の移動ができるかどうかを出力することができました。