[알고리즘] Programmers 모두 0으로 만들기 #Python

[Programmers] 모두 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	edges	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 해야 합니다.

---

📍 풀이

• 위 그림은 입출력 예 #1 에서 모든 정점의 가중치를 0으로 만드는 과정의 순서를 표현한 그림이다.
  • 2 → 4 → 3 → 0 → 1 의 순서에 따라 진행되었다.
• 과정의 순서는 헤드 노드 1 을 기준으로 후위 순회(post order) 의 순서가 적용된다는 것을 알 수 있다.
• 후위 순회를 활용하여 제일 깊은 정점부터 헤드 노드까지 모든 정점의 가중치를 0으로 만드는 과정을 진행한다.
• 하위 정점의 가중치를 모두 상위 정점의 가중치에 더하는 방식으로 진행한다.
• 진행하면서 발생한 가중치의 누적 이동값을 구한다.

---

📌 문제 풀이

✏️ [0] 후위 순회 함수

def postOrder(n):
    global answer
    for i in graph[n]:
        if visited[i]:
            visited[i] = False
            a[n] += postOrder(i)  
    answer += abs(a[n])
    return a[n]

• 현재 정점(n)에 연결된 모든 정점(i)을 재귀의 방식을 통해 방문한다.
• answer 에는 현재 정점의 가중치 값을 누적으로 더한다.

✏️ [1] 그래프 생성

graph = defaultdict(list)
headCandidate = defaultdict(int)
for n1, n2 in edges:
    graph[n1].append(n2)
    graph[n2].append(n1)
    headCandidate[n1] += 1
    headCandidate[n2] += 1

• 주어진 간선 정보(edges)를 통해 그래프(graph)를 생성한다.
• headCandidate 는 헤드 노드의 정보를 얻기 위해 선언한 자료구조이다.
  • 주어진 간선 정보(edges)를 통해 각 정점에 연결된 간선의 개수를 카운트할 수 있다.

✏️ [2] head node 찾기

for key, value in headCandidate.items():
    if value == 1:
        head = key
        break

• 정점에 연결된 간선의 개수가 1인 경우 헤드 노드(head)로 선정한다.

✏️ [3] 후위 순회

visited = [True] * len(a)
visited[head] = False
postOrder(head)

• 헤드 노드(head)를 이용하여 후위 순회 방식을 통해 그래프를 탐색한다.
• 방문한 정점의 재방문을 막기 위해 visited 리스트를 선언한다.

✍ 코드

from sys import setrecursionlimit
from collections import defaultdict

setrecursionlimit(10**6) # 재귀 깊이 제한 해제
answer = 0

def solution(a, edges):
    
    def postOrder(n):
        global answer
        for i in graph[n]:
            if visited[i]:
                visited[i] = False
                a[n] += postOrder(i)  
        answer += abs(a[n])
        return a[n]
    
    # 그래프 생성
    graph = defaultdict(list)
    headCandidate = defaultdict(int)
    for n1, n2 in edges:
        graph[n1].append(n2)
        graph[n2].append(n1)
        headCandidate[n1] += 1
        headCandidate[n2] += 1
    
    # head node 찾기
    for key, value in headCandidate.items():
        if value == 1:
            head = key
            break
    
    # post order
    visited = [True] * len(a)
    visited[head] = False
    postOrder(head)
    
    return answer if a[head] == 0 else -1