[programmers/py] 모두 0으로 만들기

모두 0으로 만들기

https://school.programmers.co.kr/learn/courses/30/lessons/76503

문제 설명

• 임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한쪽은 1 감소시킵니다.
• 주어진 행동을 통해 트리의 모든 점들의 가중치를 0으로 만드는 것이 불가능하다면 -1을, 가능하다면 최소 몇 번만에 가능한지를 찾아 return 하도록 solution 함수를 완성해주세요

문제 풀이

하나의 인덱스를 투르 노드로 잡아서 DFS로 해결
자식의 합을 부모와 계속 더해서 0이되면 문제를 풀 수 있는 것

import sys
sys.setrecursionlimit(10 ** 6)

result = 0
def solution(a, edges):
    if sum(a) != 0:
        return -1
        
    n = len(a)
    graph = [[] for i in range(n)]
    for node_a, node_b in edges:
        graph[node_a].append(node_b)
        graph[node_b].append(node_a)
        
    def dfs(child, parent, graph, a):
        global result
        for c in graph[child]:
            if c != parent:
                dfs(c, child, graph, a)
        a[parent] += a[child]
        result += abs(a[child])
        
    dfs(0, 0, graph, a)
    return result