全て0になる(Python)


1.アイデア


どうやってアプローチするのかわからなかったので、解説を見ました.正式な説明参照.aのすべての要素の合計を0とし、0または-1を返します.0の場合、edgesを使用してツリーを作成します.その後、任意のノードを最初に参照するノードとして選択し、dfsに進む.dfsから得られた結果にaを加え、絶対値(dfsから得られた結果)を算出する.
加法の交換法則により,演算順序を決める必要はないと解説されている.したがって、作成されたコードについては、最初に参照するノードとして0が選択される.

2.コード

from collections import defaultdict
import sys

sys.setrecursionlimit(int(1e6))
tree = defaultdict(list)
visited = []
arr = []
cnt = 0

def dfs(x):
    global visited, tree, arr, cnt
    if visited[x]: return 0
    visited[x] = True
    for nx in tree[x]:
        k = dfs(nx)
        cnt += abs(k)
        arr[x] += k
    
    v = arr[x]
    arr[x] = 0
    return v

def solution(a, edges):
    global visited, tree, arr, cnt
    if sum(a) != 0: return -1
    n = len(a)
    visited = [False for _ in range(n)]
    arr = a
    for a, b in edges:
        tree[a].append(b)
        tree[b].append(a)
    
    dfs(0)
    
    return cnt