import sys
input = sys.stdin.readline
sys.setrecursionlimit(1<<19)
def DFS(now):
for next in child[now]:
DFS(next)
childcnt[now] += childcnt[next]
result[now] += result[next]+cost[next]*childcnt[next]
def solve(now):
result[now] = result[parent[now]]+cost[now]*(N-childcnt[now]*2)
for next in child[now]:
solve(next)
N = int(input())
graph = [[] for i in range(N)]
for _ in range(N-1):
x,y,w = map(int,input().split())
graph[x-1].append((y-1,w)); graph[y-1].append((x-1,w))
parent = [0]*N; child = [[] for i in range(N)]; cost = [0]*N
dq = [0]
while dq:
now = dq.pop()
for next,w in graph[now]:
if child[next]:
continue
dq.append(next)
child[now].append(next)
parent[next],cost[next] = now,w
result = [0]*N; childcnt = [1]*N
DFS(0); solve(0)
print(*result,sep="\n")
트리 응용 문제
DFS를 두 번 순회해야하는 문제였다.
과정을 간략히 설명하면,
1. BFS를 통해 각 노드의 부모노드, 자식노드, 부모로 이동하는 비용을 저장하여 루트노드가 1인 트리구조를 만든다.
2. DFS를 통해 각 노드에서 (자신을 포함한 자식노드의 수) 와 (모든 자식노드로 이동하는 비용)을 구한다.
3. DFS를 한번 더 한다. 이때는 2에서 구한 모든 자식노드로 이동하는 비용을 이용하여, 모든 노드로 이동하는 비용을 구한다.
이제 각 과정을 설명하면,
1에서 트리 구조 만드는건 간단하므로 넘어가고,
2에서 자식노드의 수는 (자식노드들의 자식노드의 수)의 합이다. 또, 모든 자식노드로 이동하는 비용은, (자식노드들의 모든 자식노드로 이동하는 비용 + 자식노드의 부모노드로 이동하는 비용*자식노드의 자식노드수)의 합이다. 왜냐하면 자식노드와 부모노드 사이의 간선을 자식노드의 자식노드수만큼 이용하기 때문이다.
3에서 모든 노드로 이동하는 비용은, (부모노드의 모든 노드로 이동하는 비용) 에서 부모-자식의 간선 비용을 이동하는 횟수의 차이를 빼주면 된다.
설명이 조금 부실하다고 느낄 수 있는데, 이 문제의 경우 직접 트리 구조를 그려보고, 부모노드의 비용과 자식노드의 비용이 어떻게 달라지는지 그 관계성에 대해서 관찰하는 것을 추천한다.
오늘의 교훈) 다양한 문제를 해결하자