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

18 Mar 2022 in Algorithm on Programmers, Level3

[문제 링크]

풀이

from collections import defaultdict
import sys
sys.setrecursionlimit(30000)

answer = 0

def solution(a, edges):
    global answer

    if sum(a) != 0:
        return -1

    tree = defaultdict(list)
    for x, y in edges:
        tree[x].append(y)
        tree[y].append(x)

    visited = [0] * len(a)
    dfs(0, a, tree, visited)

    return answer


def dfs(x, a, tree, visited):
    global answer

    visited[x] = 1
    for next in tree[x]:
        if not visited[next]:
            a[x] += dfs(next, a, tree, visited)
    answer += abs(a[x])

    return a[x]

트리 구조는 어떤 노드든 루트 노드가 될 수 있다. 따라서 시작점이 어떤 곳이든지 상관 없다. 시작점을 기준으로 재귀를 돌려서 값들을 쌓아 올리면 된다.