[BOJ] 15647. 로스팅하는 엠마도 바리스타입니다

https://www.acmicpc.net/problem/15647

트리의 각 노드에서 다른 모든 노드로 가는 최단 경로들의 합을 구하는 문제입니다.

간단하게 생각해서 각 노드가 다른 노드로 갈 때 무조건 루트 노드를 거친다고 가정하고 계산한 후 필요없는 간선을 빼주면 될 거라고 생각했습니다.

len_to_root = [0 for _ in range(N + 1)] # 루트 노드에서 각 노드로 가는 거리
len_to_parent = [0 for _ in range(N + 1)] # 부모 노드로 가는 거리
parents = [0 for _ in range(N + 1)] # 부모 노드
subtree = [0 for _ in range(N + 1)] # 서브트리 개수

def dfs(now, parent, length):
    len_to_root[now] = length
    parents[now] = parent
    subtree[now] = 1
    for edge in edges[now]:
        if parent == edge[0]: # 부모 노드로 가는 간선일 경우
            len_to_parent[now] = edge[1] # 부모 노드로 가는 거리 저장
            continue
        subtree[now] += dfs(edge[0], now, length + edge[1]) # 서브트리 개수 저장
    return subtree[now]

dfs(root, -1, 0)

dfs를 돌리면서, 각 노드의 서브트리 개수와 바로 부모 노드와의 거리, 루트 노드까지의 거리를 저장합니다.

minus = [0 for _ in range(N + 1)]

def dfs2(now, m):
    minus[now] = m + len_to_parent[now] * subtree[now] * 2
    for edge in edges[now]:
        if parents[now] == edge[0]:
            continue
        dfs2(edge[0], minus[now])

dfs2(root, 0)

한번 더 dfs를 돌리면서 바로 위 부모 노드와의 거리와 서브트리 개수를 곱한 값을 자식 노드로 가면서 더해줍니다.
이 값은 루트 노드를 거치지 않는 경우에 빼줘야 하는 값입니다.

all = 0
for length in len_to_root:
    all += length

for i in range(1, N + 1):
    res = all + len_to_root[i] * N
    res -= minus[i]
    print(res)

res값은 처음에 루트 노드를 거치는 경우를 가정합니다.
그 후에 필요없는 간선을 빼주고 출력합니다.