[프로그래머스] 모두 0으로 만들기 - 파이썬

April 25, 2021

문제 설명

각 점에 가중치가 부여된 트리가 주어집니다. 당신은 다음 연산을 통하여, 이 트리의 모든 점들의 가중치를 0으로 만들고자 합니다.

임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한쪽은 1 감소시킵니다. 하지만, 모든 트리가 위의 행동을 통하여 모든 점들의 가중치를 0으로 만들 수 있는 것은 아닙니다. 당신은 주어진 트리에 대해서 해당 사항이 가능한지 판별하고, 만약 가능하다면 최소한의 행동을 통하여 모든 점들의 가중치를 0으로 만들고자 합니다.

트리의 각 점의 가중치를 의미하는 1차원 정수 배열 a와 트리의 간선 정보를 의미하는 edges가 매개변수로 주어집니다. 주어진 행동을 통해 트리의 모든 점들의 가중치를 0으로 만드는 것이 불가능하다면 -1을, 가능하다면 최소 몇 번만에 가능한지를 찾아 return 하도록 solution 함수를 완성해주세요. (만약 처음부터 트리의 모든 정점의 가중치가 0이라면, 0을 return 해야 합니다.)

제한사항

a의 길이는 2 이상 300,000 이하입니다. a의 모든 수는 각각 -1,000,000 이상 1,000,000 이하입니다. a[i]는 i번 정점의 가중치를 의미합니다. edges의 행의 개수는 (a의 길이 - 1)입니다. edges의 각 행은 [u, v] 2개의 정수로 이루어져 있으며, 이는 u번 정점과 v번 정점이 간선으로 연결되어 있음을 의미합니다. edges가 나타내는 그래프는 항상 트리로 주어집니다.

입출력 예

내 코드

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

def solution(a, edges):
    if sum(a) != 0 :
        return -1
    graph = defaultdict(list)
    for u, v in edges :
        graph[u].append(v)
        graph[v].append(u)
    def dfs(u,v) :
        global answer
        for node in graph[u] :
            if node != v :
                dfs(node,u)
        answer += abs(a[u])
        a[v] += a[u]
        a[u] = 0
    dfs(0,0)
    return answer

먼저 모든 노드의 가중치의 합이 0이 되지않으면 0으로 만들기가 불가하기 때문에 -1을 리턴처리하고, graph를 defaultdict()를 통해서 받고 dfs를 통해 루트노드에서 리프노드 까지 dfs를 호출하여 리프노드부터 값을 계산하였다. 그래프 문제는 머릿속에 그려가면서 어떤 로직으로 풀어야할지를 해야겠다.