Python3で解く AtCoder Regular Contest 106 B - Value

Union-Findを使って、グループ間の合計値が一致すれば可能。

概要

問題

$n$ 個の頂点、 $m$ 個の辺がある単純無向グラフが与えられる。
$i$ 番目の辺は頂点 $c_i$ と頂点 $d_i$ を結んでいる。
始め、頂点 $i$ には値 $a_i$ が書かれている。
以下の操作を回以上行う。
1. ひとつ辺を選んで結ばれている頂点を, とする。そのとき以下の操作のどちらかを行う。
  1. 値 $a_x$ を $-1$ し、値 $a_y$ を $+1$ する。
  2. 値 $a_x$ を $+1$ し、値 $a_y$ を $-1$ する。

解くときに考えた内容

要するに、辺を選んで、隣り合っていれば+1, -1を回数制限なしでできる。つまり、連結さえしていれば、それらの値は変更自在で、連結している要素数の合計値が一致していればOKで、そうでないならダメっていうこと。細かい証明はこちらに委ねる。

この場合どのように実装するかといえば、連結しているグループの話なので、Union-Findを用いて連結しているグループの値の和が一致するかを検討するだけ。このとき注意しなくてはいけないのは、find()でrootを求めていく。

まずは、頂点 $c_i$ と頂点 $d_i$ の情報から、どことどこが連結しているかをUnion-Findで求める。
その後、頂点 $0$ から頂点 $n-1$ までどのルートに属しているかをチェック。
（find()でもとめる、そのルートを $x$ とする）
ルート $x$ の $a[i]$ の合計とルート $x$ の $b[i]$ の合計をそれぞれ計算する。
全てのルートのぞれぞれの合計値一致すればOK

コード

# Union-Find
def find(target):
    if parent[target] < 0:
        return target
    else:
        parent[target] = find(parent[target])
        return parent[target]
def union(x, y):
    root_x = find(x)
    root_y = find(y)
    if root_x == root_y:
        return
    if parent[root_x] > parent[root_y]:
        root_x, root_y = root_y, root_x
    parent[root_x] += parent[root_y]
    parent[root_y] = root_x

n, m = map(int, input().split())
parent = [-1 for _ in range(n)]
a = [int(x) for x in input().split()]
b = [int(x) for x in input().split()]
for _ in range(m):
    c, d = map(lambda x: int(x) - 1, input().split())
    union(c, d)
# rootにつながっているグループの和の計算用
a_sum = [0 for _ in range(n)]
b_sum = [0 for _ in range(n)]
for i in range(n):
    # i番目のroot
    x = find(i)
    # i番目が属するrootにa[i]を足していく
    a_sum[x] += a[i]
    # i番目が属するrootにb[i]を足していく
    b_sum[x] += b[i]
# リストの中身が完全一致すればOK
if a_sum == b_sum:
    print('Yes')
else:
    print('No')

参考になる書籍

問題解決力を鍛える!アルゴリズムとデータ構造

けんちょんさんこと、大槻さんの書籍である、「問題解決力を鍛える!アルゴリズムとデータ構造」ご本人の記事の中で難易度は、螺旋本ということで初学者でも非常に結構読みやすい!!!（実際に、けんちょんさんの記事はC++に疎い自分でも、その思考プロセスを学べるくらい丁寧に書かれているし、この書籍も非常に読みやすい）どのようなステータスの人にもおすすめの一冊になっている。