atcoder ARC 106 B - Valuesの勉強

atcoder Union-Find

atcoder.jp

・参考
　AtCoder Regular Contest 106 B – Values を解いた記録 – Manuel1024の引きこもルーム

・説明
　Union-Findを使って実装した。操作によって連結成分毎の合計は変わらないことので、もし変わっていればNoになる。
　atcoderで新しく使えるようになったac-library-python/atcoder at master · not522/ac-library-python · GitHubを使って実装した。

from atcoder import dsu

n,m=map(int,input().split())
a=list(map(int,input().split()))
b=list(map(int,input().split()))

uf=dsu.DSU(n)

for i in range(m):
    c,d=map(int,input().split())
    c-=1
    d-=1
    uf.merge(c,d)

a_sum=[0 for _ in range(n)]  #操作前の連結成分毎の合計.
b_sum=[0 for _ in range(n)]  #操作後の連結成分毎の合計.

#連結成分毎の合計を求める.
for i in range(n):
    a_sum[uf.leader(i)]+=a[i]
    b_sum[uf.leader(i)]+=b[i]

for i in range(n):
    if a_sum[i]!=b_sum[i]:
        print("No")
        exit()
print("Yes")