ARC106 B - Values

備忘録

問題

atcoder.jp

回答

import sys
import os
import math
import bisect
import itertools
import collections
import heapq
import queue
import array
import time

# 時々使う
# from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall
# from decimal import Decimal
# from collections import defaultdict, deque

# 再帰の制限設定
sys.setrecursionlimit(10000000)


def ii(): return int(sys.stdin.buffer.readline().rstrip())
def il(): return list(map(int, sys.stdin.buffer.readline().split()))
def fl(): return list(map(float, sys.stdin.buffer.readline().split()))
def iln(n): return [int(sys.stdin.buffer.readline().rstrip())
                    for _ in range(n)]


def iss(): return sys.stdin.buffer.readline().decode().rstrip()
def sl(): return list(map(str, sys.stdin.buffer.readline().decode().split()))
def isn(n): return [sys.stdin.buffer.readline().decode().rstrip()
                    for _ in range(n)]


def lcm(x, y): return (x * y) // math.gcd(x, y)


MOD = 10 ** 9 + 7
INF = float('inf')


class UnionFind:
    def __init__(self, N):
        self.tree = [-1 for _ in range(N)]

    def root(self, N):
        if self.tree[N] < 0:
            return N
        self.tree[N] = self.root(self.tree[N])
        return self.tree[N]

    def same(self, x, y):
        return self.root(x) == self.root(y)

    def unite(self, x, y):
        x = self.root(x)
        y = self.root(y)
        if x == y:
            return
        if self.tree[x] > self.tree[y]:
            x, y = y, x
        self.tree[x] += self.tree[y]
        self.tree[y] = x

    def size(self, x):
        return -self.tree[self.root(x)]


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N, M = il()
    A = il()
    B = il()

    uf = UnionFind(N)
    for _ in range(M):
        c, d = il()
        uf.unite(c-1, d-1)

    ret = [[0, 0] for _ in range(N)]
    for n in range(N):
        p = uf.root(n)
        ret[p][0] += A[n]
        ret[p][1] += B[n]

    for a, b in ret:
        if a != b:
            print('No')
            break
    else:
        print('Yes')


if __name__ == '__main__':
    main()

考え方

指示された操作を何度行っても、総和が不変であることに注目します

• 回答概要
  • UnionFind木で連結成分を求める
  • 連結成分の総和が等しいか確かめる

---

ある辺を選択して、指示された操作を行います。
このとき、辺の両端に存在する頂点の総和は常に不変です。
同様に、操作を何度行っても、１つの連結成分の総和は不変です。

各頂点の値は１つの連結成分の総和以上にすることはできません。
ただし、どこに位置する頂点でも連結成分の総和以下の値には必ず操作することが出来ます。

以上のことから、
各a1, a2, ... , anの連結成分の総和と、各b1, b2, ... ,bnの連結成分の総和が等しいとき、
操作によって各aと各bを等しくすることが出来ます。

実装はUnionFind木を使用して、連結成分を探索しました。
連結成分が判明したのち、各連結成分ごとの総和を求めて、
比較を行うことで目的を達成することが可能か判定しました。