unionfindとdefalutdict

unionfindの勉強をしたので。
unionfindはグラフの頂点で互いに繋がっているかどうかを判断するのに便利です。
追加の仕方がややこしくて、まだ全然、体に馴染んでない。
unionfindのプログラムの肝は、繋がっている頂点に対しては、同じ値を返すという関数です。
今回は、findという名前で実装しています。

解いた問題はABC049のD問題。
D: 連結 / Connectivity - AtCoder Beginner Contest 049 | AtCoder
解説はこれ。ほとんど何も書いてない。
https://atcoder.jp/img/arc065/editorial.pdf

この問題をやってるときに、初め、Counterで数え上げを行おうと思っていた。
しかし、全然うまく行かなかった。もちろん、うまく実装してる人もいたが。
なぜうまく行かないか、というと、自分のプログラムでできるunion-find treeが必ずしも最適なものとは限らないからだった。
つまり、複数の再帰を必要とするため、Counterとの相性が良くなかった。

というわけで、defaultdictを用いて、実装した。
参考文献はこれ。
8.3. collections — コンテナデータ型 — Python 3.6.5 ドキュメント
Python defaultdict の使い方

pythonの公式ドキュメントは日本語なので神。

プログラムはこれ。

from collections import defaultdict
def find(x, A): #unionfind
    p = A[x]
    if p == x:#根になっていればよい
        return x
    a = find(p, A)#根はない場合は根に到達するまで再帰的にfindする。この1行で少し最適化される
    A[x] = a
    return a

N, K, L = map( int, input().split())
V = [ i for i in range(N)]#道路
W = [ i for i in range(N)]#鉄道
for _ in range(K):#道路
    p, q = map( int, input().split())
    p, q = p-1, q-1
    bp, bq = find(p,V), find(q,V)
    V[q] = bp#pの根bpにqとqの根bpをくっつける
    V[bq] = bp

for _ in range(L):#鉄道
    r, s = map( int, input().split())
    r, s = r-1, s-1
    br, bs = find(r,W), find(s,W)
    W[s] = br#rの根brにsとsの根bsをくっつける
    W[bs] = br

d = defaultdict(int) #デフォルトで0を返すdict。
for i in range(N):
    d[(find(i,V), find(i, W))] += 1
ANS = [0]*N

for i in range(N):
    ANS[i] = d[(find(i,V), find(i, W))]

print(' '.join( map( str, ANS)))

おまけ

unionfindは逐次的な操作との相性が良い。
というか、そうでない操作との相性が悪い。
例えば、defalutdictとの相性は良いが、Counterとの相性はよくない。