문제 요약
각 정점에 정수 가중치가 부여된 트리가 주어진다.
연결된 두 정점을 선택하여 한쪽 가중치는 1 증가시키고, 다른 쪽 가중치는 1 감소시키는 연산을 수행할 수 있다.
모든 정점의 가중치를 0으로 만들 수 없다면 -1, 가능하다면 필요한 최소 연산 횟수를 반환해야 한다.
제한사항에서 확인할 점
- 정점 수는 최대
300,000개이다. - 트리의 간선 수는 항상
정점 수 - 1개이다. - 가중치의 절댓값은 최대
1,000,000이다. - 트리의 깊이는 최대
300,000이 될 수 있다.
따라서 재귀 DFS를 사용하면 파이썬의 재귀 깊이 제한을 초과할 수 있다. 스택을 이용한 반복형 DFS로 구현하는 것이 안전하다.
모든 가중치를 0으로 만들 수 있는 조건
한 번 연산할 때 한 정점은 1 증가하고 다른 정점은 1 감소한다.
따라서 연산 전후의 전체 가중치 합은 항상 같다.
모든 가중치가 0인 상태의 전체 합은 0이므로 다음 조건을 만족해야 한다.
sum(a) == 0전체 합이 0이 아니면 어떤 연산을 수행하더라도 모든 정점을 0으로 만들 수 없다.
반대로 트리는 모든 정점이 연결되어 있으므로 전체 합이 0이라면 리프부터 가중치를 부모에게 전달하여 모든 정점을 0으로 만들 수 있다.
핵심 아이디어
임의의 정점을 루트로 정한 뒤 자식 정점부터 부모 정점 방향으로 가중치를 전달한다.
현재 정점을 루트로 하는 서브트리의 가중치 합이 S라고 하자.
이 서브트리와 나머지 트리를 연결하는 간선은 부모와 연결된 간선 하나뿐이다. 서브트리의 가중치를 모두 0으로 만들려면 이 간선을 통해 반드시 S만큼의 가중치를 이동해야 한다.
필요한 연산 횟수는 방향과 관계없이 다음과 같다.
abs(S)따라서 리프부터 다음 작업을 반복한다.
- 현재 정점의 누적 가중치 절댓값을 정답에 더한다.
- 현재 정점의 누적 가중치를 부모에게 더한다.
- 현재 정점은 처리된 것으로 본다.
최소 연산 횟수가 되는 이유
서브트리의 가중치 합이 S라면 부모 간선을 통해 최소 abs(S)만큼 이동해야 한다.
한 번의 연산으로 간선을 따라 이동할 수 있는 가중치는 1이므로 해당 간선에서 최소 abs(S)번의 연산이 필요하다.
후위 순회 방식은 각 간선에서 정확히 abs(S)번만 수행한다. 모든 간선에서 필요한 최소 횟수만 사용하므로 전체 연산 횟수도 최소가 된다.
반복형 후위 순회 만들기
반복형 DFS로 루트부터 정점을 방문하면서 다음 정보를 저장한다.
parent[node]: 현재 정점의 부모order: DFS 방문 순서
DFS 방문 순서는 부모가 자식보다 먼저 들어간다. 따라서 order를 뒤집으면 자식이 부모보다 먼저 처리되는 후위 순서를 얻을 수 있다.
DFS 방문 순서: 루트 → 부모 → 자식
역순 처리: 자식 → 부모 → 루트Python 코드
def solution(a, edges):
if sum(a) != 0:
return -1
n = len(a)
graph = [[] for _ in range(n)]
for u, v in edges:
graph[u].append(v)
graph[v].append(u)
# 원본 입력을 변경하지 않기 위해 복사한다.
weights = a[:]
# -2는 아직 방문하지 않은 상태를 의미한다.
parent = [-2] * n
parent[0] = -1
order = []
stack = [0]
while stack:
node = stack.pop()
order.append(node)
for next_node in graph[node]:
if parent[next_node] != -2:
continue
parent[next_node] = node
stack.append(next_node)
answer = 0
# 루트는 부모 간선이 없으므로 제외한다.
for node in reversed(order[1:]):
answer += abs(weights[node])
weights[parent[node]] += weights[node]
return answer코드 동작 설명
1. 전체 가중치 합 검사
if sum(a) != 0:
return -1전체 가중치 합은 연산으로 바꿀 수 없다. 합이 0이 아니면 즉시 -1을 반환한다.
2. 트리의 인접 리스트 생성
graph = [[] for _ in range(n)]
for u, v in edges:
graph[u].append(v)
graph[v].append(u)간선은 양방향이므로 두 정점의 인접 리스트에 서로를 추가한다.
3. 부모 관계와 방문 순서 기록
parent = [-2] * n
parent[0] = -1
order = []
stack = [0]parent가 -2인 정점은 아직 방문하지 않은 정점이다. 0번 정점을 루트로 정하고 반복형 DFS를 수행한다.
트리는 연결되어 있으므로 DFS가 끝나면 모든 정점이 order에 한 번씩 저장된다.
4. 자식의 가중치를 부모에게 전달
for node in reversed(order[1:]):
answer += abs(weights[node])
weights[parent[node]] += weights[node]역순으로 처리할 때 weights[node]에는 현재 정점의 모든 자손을 포함한 서브트리 가중치 합이 저장되어 있다.
이 값을 부모에게 전달하면서 필요한 최소 연산 횟수인 abs(weights[node])를 누적한다.
order[1:]을 사용하는 이유는 루트에는 부모로 연결된 간선이 없기 때문이다.
정확성
전체 가중치 합이 0이 아니면 연산으로 전체 합을 변경할 수 없으므로 모든 가중치를 0으로 만들 수 없다.
전체 합이 0인 경우, 역순 방문에서 모든 자식은 부모보다 먼저 처리된다. 따라서 현재 정점의 가중치는 해당 정점을 루트로 하는 서브트리의 가중치 합이다.
서브트리와 나머지 트리를 연결하는 간선은 부모 간선 하나뿐이므로 서브트리 합의 절댓값만큼 반드시 이 간선을 통해 이동해야 한다. 알고리즘은 정확히 그 횟수를 정답에 더하고 가중치를 부모에게 전달한다.
모든 자식을 처리한 후에는 루트에 전체 가중치 합이 모인다. 전체 합이 0이므로 루트의 최종 가중치도 0이 된다.
따라서 알고리즘은 모든 정점의 가중치를 0으로 만들며, 각 간선에서 필요한 최소 연산만 수행하므로 전체 연산 횟수도 최소이다.
시간 복잡도
정점 수를 V, 간선 수를 E라고 하자.
인접 리스트 생성, DFS, 역순 처리는 각각 모든 정점과 간선을 한 번씩 확인한다.
트리에서는 E = V - 1이므로 시간 복잡도는 다음과 같다.
O(V + E) = O(V)공간 복잡도
인접 리스트, 부모 배열, 방문 순서와 DFS 스택을 사용한다.
O(V + E) = O(V)주의할 점
- 전체 합이
0인지 먼저 확인해야 한다. - 정점이 최대 30만 개이므로 재귀 DFS보다 반복형 DFS가 안전하다.
- 부모 배열을 방문 여부로 함께 사용하면 중복 방문을 방지할 수 있다.
- 원본 배열
a를 유지하려면 복사본에서 누적해야 한다. - 정답이 매우 커질 수 있지만 파이썬 정수는 크기 제한 없이 사용할 수 있다.
- 처음부터 모든 가중치가
0이면 각 정점에서 더해지는 값도0이므로 결과는0이다.
정리
이 문제의 핵심은 각 서브트리의 가중치 합을 부모에게 전달하는 것이다.
전체 합이 0인지 확인하고, 반복형 DFS의 방문 순서를 뒤집어 자식부터 처리하면서 각 누적 가중치의 절댓값을 더하면 최소 연산 횟수를 O(V)에 구할 수 있다.
