백준 1167번 트리의 지름

백준 1167번: 트리의 지름

아이디어

트리에서 아무 노드나 하나 잡고(루트 노드) dfs로 가장 멀리 있는 노드를 찾는다. 그리고 찾은 노드에서 가장 멀리 떨어져 있는 노드까지의 거리가 트리의 지름이다. 왜 이렇게 되는걸까? 맨 처음에는 루트에서 멀리 있는 노드 두 개 찾아서 더하면 정답인 줄 알았다. 그런데 맨 처음 잡은 노드(루트 노드)를 지나지 않으면서 지름이 되는 경우도 존재한다. 간단한 그림으로 해당 case에 대해 알아보자.

그림 한 장으로 깔끔하게 정리가 되었다.

그렇다면 루트 노드도, 루트에서 가장 멀리 떨어진 노드도 지나지 않는 경우는 없는걸까? 이 경우도 그림을 통해 알아보자.

2번 경로가 우리가 찾는 경로고, 3번 경로가 루트 노드, 루트 노드와 가장 멀리 떨어진 노드 둘 다 지나지 않는 경로다. 부등식을 세워서 조금만 계산해 보면, 3번 경로는 2번 경로보다 길 수 없다는 사실을 알 수 있다.

코드

import sys
input = sys.stdin.readline

n = int(input())
tree = [[] for _ in range(n+1)]
dist = [0] * (n+1)
visited = [False] * (n+1)

for _ in range(n):
    input_list = list(map(int, input().split()))
    p = input_list[0]
    for i in range(1, len(input_list) - 1, 2):
        if i < 0:
            break
        tree[p].append((input_list[i], input_list[i+1]))

def dfs(n, v):
    for i in tree[n]:
        if not visited[i[0]]:
            visited[i[0]] = True
            d = dfs(i[0], v+i[1])
            dist[i[0]] = d
    return v

visited[1] = True
dfs(1, 0)

m = 0
idx = 0
for i in range(1, n+1):
    if dist[i] > m:
        m = dist[i]
        idx = i

visited = [False] * (n+1)
dist = [0] * (n+1)
visited[idx] = True
dfs(idx, 0)

print(max(dist))

여담

이런걸 그냥 바로바로 생각해내는 사람들은 머리가 얼마 나 좋은거지? 난 금붕어다.