[Softeer] 거리 합 구하기

문제

현호는 사내 네트워크 분석 업무를 담당하게 되었다. 현재 사내 네트워크는 N개의 노드를 가지는 트리 형태의 네트워크인데, 이 말은 두 노드간의 연결이 정확히 N-1개 있어서 이 연결만으로 모든 노드간에 통신을 할 수 있다는 뜻이다.

 

각 노드에 1에서 N사이의 번호를 붙이면 i번째 연결은 xi번 노드와 yi번 노드를 양방향으로 연결하며, 통신에 걸리는 시간은 ti이다. D(i,j)는 i번 노드와 j번 노드 사이의 거리를 나타내는데, i번 노드에서 여러 연결을 거쳐 j번 노드에 도달하기 위해 걸리는 최소 시간이다.

 

노드를 들를 때 추가적인 작업이 없는 이상적인 시간을 따진다. 현호는 네트워크 분석을 위해 어떤 노드 i를 기준으로 다른 모든 노드 사이와의 거리의 합을 알고 싶다. 즉, 

을 알고 싶다.

 

입력예제 2번을 예로 들면, 다음과 같이 7개의 노드로 이루어진 네트워크로 표현할 수 있다. 각 연결 위에 적힌 수는 t를 나타낸다.

 

 

1번 노드를 기준으로 D(1,j)를 구해보면 다음과 같다.

 

제약조건

1≤N≤2×105

주어진 N-1개의 연결로 모든 노드가 연결되어 있다.

1≤xi,yi≤N

1≤ti≤106

입력형식

첫 번째 줄에 노드의 개수 N이 주어진다.

다음 N-1줄의 i번째 줄에는 i번째 연결을 나타내는 세 정수 xi,yi,ti가 주어진다.

출력형식

N개의 줄에 걸쳐서, i번째 줄에는 i번 노드와 다른 모든 노드 사이의 거리의 합, 

를 출력한다.

 

입력예제1

4 1 2 1 2 3 2 3 4 4

출력예제1

11 9 9 17

입력예제2

7 1 2 5 1 3 2 1 4 8 3 5 4 3 6 1 4 7 6

출력예제2

38 63 40 62 60 45 92

---

풀이

풀이1(오답)

import sys
from collections import deque
input = sys.stdin.readline

n = int(input())
node = [[] for _ in range(n+1)]
distance = [[0] * (n+1) for _ in range(n+1)]

for _ in range(n-1):
    x, y, t = map(int, input().split())
    node[x].append(y)
    node[y].append(x)
    distance[x][y], distance[y][x] = t, t

q = deque([])

for i in range(1, n+1):
    q.append((i,i,0))
    visited = [0] * (n+1)
    visited[i] = 1
    
    while q:
        start, now, d = q.pop()
        
        for y in node[now]:
            # print(y)
            # print("-------")
            
            if visited[y] == 0:
                visited[y] = 1
                distance[start][y] = d+distance[now][y]
                q.append((start,y,distance[start][y]))

for i in range(1, n+1):
    print(sum(distance[i]))

 

 

예제 테케에 대해서는 맞지만 메모리 초과로 틀리는 방법인듯. 제출하면 아예 이력이 안뜬다 ㅠㅜ

 

풀이2

 

이 문제에서 중요한 부분은 서브트리를 만드는 것이었던 것 같다.

서브트리란

부모노드와의 간선을 잘랐을 때, 아래로 존재하는 노드들의 개수이다.

최상위 노드는 부모 노드가 없기 때문에 모든 노드의 개수이고, 리프 노드(최하위 노드)들은 자식 노드가 없기 때문에 자기 자신 밖에 없다.

리프노드는 자기자신 밖에 없기 때문에 최하위부터 더해가는 식으로 BottomUp 방식으로 리프노드의 개수를 구해야 한다.

 

1번에서 각 노드까지 길이의 총합은 각 노드로의 길이를 더하면 되지만,

예를 들어, 2번에서 향할 경우에는, 서브트리에 있는 5,6번을 제외하고는 2번에서 1번까지의 길이를 더해주어야 한다.

 

위 내용을 코드로 나타내면 아래와 같다.

 

import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline

n = int(input())
node = [[] for _ in range(n+1)]
for _ in range(n-1):
    x, y, t = map(int, input().split())
    node[x].append((y, t))
    node[y].append((x, t))

subTreeSize = [0] * (n+1)
distSum = [0] * (n+1)

def bottomUp(now, parent):
    subTreeSize[now] = 1
    for child, weight in node[now]:
        if child != parent:
            bottomUp(child, now)
            distSum[now] += distSum[child] + subTreeSize[child]*weight
            subTreeSize[now] += subTreeSize[child]
    return

def topDown(now, parent):
    for child, weight in node[now]:
        if child != parent:
            distSum[child] = distSum[now] + weight*(n-2*subTreeSize[child])
            topDown(child, now)
    return

bottomUp(1,1)
topDown(1,1)

for i in range(1, n+1):
    print(distSum[i])