[Softeer] 거리합 구하기

https://softeer.ai/practice/info.do?idx=1&eid=635&sw_prbl_sbms_sn=148501

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

​

코드

import sys
sys.setrecursionlimit(10**6)

def dfs(current, parent):
    subtreeSize[current] = 1
    for i in range(len(graph[current])):
        child = graph[current][i][0]
        weight = graph[current][i][1]
        if child != parent:
            dfs(child, current)
            distSum[current] += distSum[child] + subtreeSize[child]*weight
            subtreeSize[current] += subtreeSize[child]
    return

def dfs2(current, parent):
    for i in range(len(graph[current])):
        child = graph[current][i][0]
        weight = graph[current][i][1]
        if child != parent:
            distSum[child] = distSum[current] + weight*(n-2*subtreeSize[child])
            dfs2(child, current)
    return

n = int(input())
graph = [[] for _ in range(n+1)]
distSum = [0]*(n+1)
subtreeSize = [0]*(n+1)

for _ in range(n-1):
    a, b, cost = map(int, input().split())
    graph[a].append([b, cost])
    graph[b].append([a, cost])

dfs(1, 1)
dfs2(1, 1)

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

맨 처음 단순히 dfs 를 사용해서 각 위치에서의 거리합들을 구하였다. 그렇게 했더니 O(N2)이 되어서 시간초과가 나왔다. 이번 문제의 경우 Softeer에서 해설영상이 제공된다.

​

https://softeer.ai/community/view.do?idx=674&cd=edu&pageNo=1

Softeer