[프로그래머스] 모두 0으로 만들기(Python)

1. 아이디어

어떻게 접근할 지 몰라서 해설을 봤다. 공식 해설을 참조했다. a의 모든 원소의 합이 0이되는 지 확인한다음, 0이 아니면 -1를 리턴하고 0이면 edges를 이용하여 트리를 만든다. 이후, 임의의 노드를 1번째로 탐색할 노드로 선택해서 dfs한다. dfs에서 나온 결과를 a에 더하고, 절대 값을(dfs에서 나온 결과)을 카운트한다.

해설에 따르면 덧셈의 교환법칙 때문에, 연산 순서를 정할 필요가 없다고한다. 따라서 작성된 코드의 경우, 임의의 노드로 0을 1번째로 탐색할 노드로 선택했다.

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