【AtCoder版！蟻本】ATC001 B - Union Find【Union-Find】

AtCoder版！蟻本のUnion-Find木の例題としてあげられているATC001 B - Union FindをPythonで解いていきます。

リンク

B - Union Find

ATC001 B - Union Find

問題文

N 頂点の、単純とは限らない無向グラフを考えます。 初期状態では、頂点のみが存在し、辺は全く存在せず、全ての頂点が孤しているとします。 以下の 2 種類のクエリが、Q 回与えられます。

・連結クエリ： 頂点 A と、頂点 B を繋ぐ辺を追加します。
・判定クエリ： 頂点 A と、頂点 B が、連結であるかどうか判定します。連結であれば Yes、そうでなければ No を出力します。
クエリを順番に処理し、判定クエリへの回答を出力して下さい。 この際、同じ辺が何度も追加されることや、自分自身への辺が追加されることもある事に注意してください。

連結であるとは、頂点 A から頂点 B まで辺をたどって到達可能であることを意味します。
A と B が同じ頂点の場合、連結であると定義します。 グラフは無向であるため、連結クエリによって頂点 A,B 間の辺が追加されると、A から B へも B から A へも辿れるようになります。

制約

・1 行目には、頂点の個数を表す整数 N(1≦N≦100,000) と、クエリの個数を表す整数 Q(1≦Q≦200,000) が、スペース区切りで与えられる。
・2 行目からの Q 行のうち i 行目には、i 番目のクエリの種類を表す整数 Pi(0≦Pi≦1) と、クエリの対象となる頂点の番号を表す整数 Ai,Bi(0≦Ai,Bi≦N−1) が、 スペース区切りで与えられる。
・Pi=0 の時、そのクエリが連結クエリであることを表す。
・Pi=1 の時、そのクエリが判定クエリであることを表す

解答例

max_num = int(1e7)
# ノードiの親がpar[i]
par = [0 for _ in range(max_num)]
# 木の深さ
rank = [0 for _ in range(max_num)]

# 初期化
def init(n):
    for i in range(n):
        # 初期化では自身を親とする
        par[i] = i
        # rankは0
        rank[i]= 0

# 木の根を求める
def find(x):
    # 自身が親なら自身の番号を返す
    if par[x] == x:
        return x
    else:
        # 根を張り替える
        par[x] = find(par[x])
        return par[x]

# xとyの属する集合を併合
def unite(x,y):
    root_x = find(x)
    root_y = find(y)
    # 同じグループの場合
    if root_x == root_y:
        return
    # ランクの大きいほうの根に小さいほうの根をつなぐ
    if rank[root_x] < rank[root_y]:
        par[root_x] = root_y
    else:
        par[root_y] = root_x
        # ランクが同じ場合は1だけ増加
        if rank[x] == rank[y]:
            rank[x] += 1

# xとyが同じ集団か
def same(x,y):
    return find(x) == find(y)


n, q = map(int,input().split())

init(n)

for i in range(q):
    p, a, b = map(int, input().split())
    # pが0ならaとbを併合
    if p == 0:
        unite(a, b)

    # pが1なら判定
    else:
        if same(a, b):
            print("Yes")
        else:
            print("No")

解説

ATC001 Bのページに解説があり、そちらを参考にしたほうが良いとは思いますが、一応記事を書いていきます。

Union-Findの練習問題で、pが0ならaとbを連結し、pが1ならaとbが連結かどうかを答える問題です。

pが0のとき、aとbをunite関数で連結させます。
このとき、根の深さを比べて、根が深いほうに浅い集合を連結させることで必要以上に根が深くなることを防ぎます。

pが1のときは、aとbの根をfind関数で調べます。
根が同じであれば、同じ集団に属しており連結であることがわかります。

この判定結果に従って解答することでOKです。

AtCoder版！蟻本（初級編）の記事リンクページはこちら
ebisuke33.hatenablog.com

リンク