Softeer Level 3 거리 합 구하기 Python

문제 링크

https://softeer.ai/practice/6258

Softeer - 현대자동차그룹 SW인재확보플랫폼

정답 코드

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

nodeDict = dict()

N = int(input())
adjL = [[] for _ in range(N+1)]
for _ in range(N-1):
    start, to, distance = map(int, input().split())
    adjL[start].append(to)
    adjL[to].append(start)
    nodeDict[(start, to)] = distance
    nodeDict[(to, start)] = distance

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

def bottomUp(childNode, parentNode):
    subTreeSize[childNode] = 1
    for nextChildNode in adjL[childNode]:
        if nextChildNode == parentNode:
            continue
        bottomUp(nextChildNode, childNode)
        totalDists[childNode] += totalDists[nextChildNode] + subTreeSize[nextChildNode] * nodeDict[(childNode, nextChildNode)]
        subTreeSize[childNode] += subTreeSize[nextChildNode]
    return

def topDown(childNode, parentNode):
    for nextChildNode in adjL[childNode]:
        if nextChildNode == parentNode:
            continue
        totalDists[nextChildNode] = totalDists[childNode] + (N - 2 * subTreeSize[nextChildNode]) * nodeDict[(childNode, nextChildNode)]
        topDown(nextChildNode, childNode)
    return


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

for idx in range(1, N+1):
    print(totalDists[idx])

풀이 방법

1. 트리 형식으로 거리가 중복 되는 현상을 고려한다.
2. 아래에서 위로 올라가면서 거리를 저장한다.
   1. 가장 깊은 곳에 위치한 정점 노드로 이동한다.
   2. 자식 노드의 개수 subTreeSize를 더해준다.
      1. 쉽게 말해서 자식의 자식의 자식까지 subTreeSize에 저장한다.
   3. 해당 사이즈 만큼 거리를 곱하고 다음 거리에 현재 거리를 더해준다.
3. 위에서 아래로 내려가면서 거리를 저장한다.
   1. 가장 깊은 곳으로 내려가지 않고 내려간 위치에서 모든 자녀(자식의 자식의 자식)을 곱해서 현재 위치에 저장
4. totalDist를 출력한다.

느낀점

• 처음엔 이론상으로 이해가 되었는데 나중에는 제대로 이해했다.
• 무조건 bfs가 좋다고 생각했는데 dfs를 써야 할 때도 있다는 것을 알았다.