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

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

코드

import sys
sys.setrecursionlimit(300000)
from collections import defaultdict
answer = 0

def dfs(x, a, tree, visited):
    global answer
    
    visited[x] = 1

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

def solution(a, edges):
    global answer
    
    if sum(a) != 0:
        return -1
    
    tree = defaultdict(list)
    
    for i, j in edges:
        tree[i].append(j)
        tree[j].append(i)

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

풀이

각 노드들의 가중치의 합이 0이 아니라면 모든 노드의 가중치를 0으로 만들 수 있는 방법이 없으므로 -1을 반환한다.

각 노드들의 가중치의 합이 0이라면 주어진 그래프가 트리 구조라는 것을 이용해 리프 노드에서부터 루트 노드로 올라가며 연산을 진행했을 때 연산의 최소 횟수를 구할 수 있다.

이때 트리 구조는 어떤 노드든 루트 노드가 될 수 있다는 성질을 이용해 아무 노드를 루트 노드로 잡고 탐색을 진행한다. (단 모든 트리는 루트 노드를 하나만 가질 수 있다.)