Programmers - 모두 0으로 만들기[파이썬(python)]

문제 출처 : https://programmers.co.kr/learn/courses/30/lessons/76503

코딩테스트 연습 - 모두 0으로 만들기

문제 설명

각 점에 가중치가 부여된 트리가 주어집니다. 당신은 다음 연산을 통하여, 이 트리의 모든 점들의 가중치를 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가 나타내는 그래프는 항상 트리로 주어집니다.

 

입출력 예

a	adges	result
[-5,0,2,1,2]	[[0,1],[3,4],[2,3],[0,3]]	9
[0,1,0]	[[0,1],[1,2]]	-1

 

입출력 예 설명

입출력 예 #1

• 다음 그림은 주어진 트리의 모든 정점의 가중치를 0으로 만드는 과정을 나타낸 것입니다.

1. 2번 정점과 3번 정점을 선택하여 2번 정점은 1 감소시키고, 3번 정점은 1 증가시킵니다. (2번 반복)
2. 3번 정점과 4번 정점을 선택하여 4번 정점은 1 감소시키고, 3번 정점은 1 증가시킵니다. (2번 반복)
3. 0번 정점과 3번 정점을 선택하여 3번 정점은 1 감소시키고, 0번 정점은 1 증가시킵니다. (5번 반복)

• 모든 정점의 가중치를 0으로 만드는 데 필요한 최소 행동 횟수는 9번이므로, 9를 return 해야 합니다.

입출력 예 #2

• 주어진 트리는 모든 정점의 가중치를 0으로 만드는 것이 불가능하므로, -1을 return 해야 합니다.

 

---

문제 풀이

이 문제에서 입력 edge가 항상 트리 구조로 주어진다는 점에 주목해야 한다. 트리는 노드가 항상 하나의 부모만을 가지고 자식은 여러 개 가질 수 있다는 특징이 있다. 또한 위의 예시를 자세히 보면 노드가 가진 가중치의 총합이 0일 경우엔 모든 가중치를 0으로 만드는 게 가능하지만 0이 아닐 경우 모든 가중치를 0으로 만드는 게 불가능하다는 것을 알 수 있다.

 

이 두가지 특징을 두고 잘 생각해보면 노드의 모든 가중치의 합이 0일 때 트리의 리프 노트부터 루트 노드까지 거슬러 올라가며 자식의 가중치를 부모 노드에 더해주는 과정을 반복하면 루트 노드의 가중치는 최종적으로 0이 될 수 있다는 특징을 발견할 수 있다.

 

그렇다면 루트노드는 어떻게 정할 수 있을까?

사실 루트노드는 어떤 노드다라고 특정 짓지 않고 모든 노드를 리프로 두었을 때 동일한 결과를 얻을 수 있다. 그 이유는 합연산은 숫자의 순서가 바뀌어도 항상 값이 동일하기 때문이다. 하지만 제한 사항의 a는 항상 2보다 크므로 0번 노드는 어떠한 경우에도 항상 존재한다.

따라서 0번 노드를 리프노드라고 가정하고 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