백준 1167번 : 트리의 지름 in Python
이 문제는 트리의 지름, 즉 임의의 두 점 사이의 거리 중 가장 긴 것을 구하는 문제이다.
이는 단순히 각 노드마다 그 노드로부터 다른 노드들까지의 거리 중 최댓값을 구하고 이렇게 나온 각 노드에 대응하는 최댓값들의 최댓값을 구하면 쉽게 풀 수 있으나, 문제는 이렇게 알고리즘을 구현하면 V^2의 시간이 걸려 시간 초과가 발생한다.
이를 해결하기 위해 우선 임의의 노드를 잡아 그 노드로부터 가장 거리가 먼 노드를 찾는다. 그러면 그 노드는 놀랍게도 트리의 지름에 대응하는 경로의 양 끝 점 중 하나가 된다. 그러므로 이 노드에 대해 가장 거리가 먼 노드를 찾아 그 노드와의 거리를 구하면 그 값이 트리의 지름과 같다.
그렇다면 왜 임의의 노드에 대해 가장 거리가 먼 노드가 항상 트리의 지름에 대응하는 경로의 양 끝 점 중 하나가 될까? 만약 임의의 노드 x로부터 가장 거리가 먼 노드를 y라고 하자. 이때 실제 트리의 지름에 대응하는 경로의 양 끝 점을 각각 u, v라고 하면, 만약 x 또는 y가 u 또는 v가 된다면 x, y 사이의 거리가 곧 트리의 지름이 될 것이다. 그러므로 x, y, u, v가 서로 다르다고 가정해보자. 그러면 크게 x, y 사이의 경로와 u, v 사이의 경로가 적어도 한 노드에서 만나는 경우와, x, y 사이의 경로와 u, v 사이의 경로가 만나지 않는 경우로 나눌 수 있다.
우선 x, y 사이의 경로와 u, v 사이의 경로가 노드 t에서 만난다고 가정하자. 그러면 x, y 사이의 경로가 x, u 사이의 경로와 x, v 사이의 경로보다 더 크므로, t, y 사이의 경로가 t, u 사이의 경로와 t, v 사이의 경로보다 더 클 것이다. 그러면 만약에 u, v 사이의 경로에서 t, v 사이의 경로를 t, y 사이의 경로로 바꾸면 경로의 길이가 더 길어지는데, 이는 u, v 사이의 경로가 트리 내에서 가장 길다는 사실에 모순이 된다. 따라서 첫번째 경우는 성립할 수 없다.
두 번째 경우로 x, y 사이의 경로와 u, v 사이의 경로가 만나지 않는다고 가정하자. 그러면 이 그래프가 트리이기 때문에 x, y 사이의 경로와 u, v 사이의 경로를 잇는 경로가 딱 하나 있을텐데, 이때 이 경로와 x, y 사이의 경로가 만나는 노드를 t라고 하면 x, y 사이의 거리가 x, u 사이의 거리와 x, v 사이의 거리보다 크므로 t, y 사이의 거리가 t, u 사이의 거리와 t, v 사이의 거리보다 더 클 것이다. 따라서 u, v 사이의 경로에서 t, u 사이의 경로를 t, y 사이의 경로로 바꾸면 경로의 길이가 더 커진다. 이는 u, v 사이의 경로가 트리 내에서 가장 길다는 사실에 모순이 된다. 따라서 두번째 경우도 성립할 수 없다.
그러므로 임의의 노드에 대해 가장 거리가 먼 노드가 항상 트리의 지름에 대응하는 경로의 양 끝 점 중 하나가 된다. 따라서 우선 아무 노드나 잡아서 BFS 또는 DFS를 실행해 가장 거리가 먼 노드를 찾고, 그 노드에 대해 다시 한 번 BFS를 실행하면 트리의 거리를 V 시간 안에 찾을 수 있다.
이를 코드로 구현하면 다음과 같다.
from collections import deque
import sys
input = sys.stdin.readline
N = int(input())
graph = [[] for _ in range(N+1)]
for i in range(N):
arr = list(map(int, input().split()))
for j in range(1, len(arr)-2, 2):
graph[arr[0]].append([arr[j], arr[j+1]])
def BFS(n):
global N
queue = deque()
# n 노드를 시작점으로 하는 경로의 길이를 저장하는 배열
distance = [-1 for _ in range(N+1)]
distance[n] = 0
queue.append(n)
# 가장 긴 경로의 끝 노드와 경로의 길이
_max = [0, 0]
while len(queue) > 0:
u = queue.popleft()
for v, d in graph[u]:
if distance[v] == -1:
distance[v] = distance[u] + d
queue.append(v)
if _max[1] < distance[v]:
_max = [v, distance[v]]
return _max
# 우선 가장 긴 경로의 양 끝 노드 중 하나를 찾는다.
k, distance = BFS(1)
# 그 다음에 그 끝 노드에 대해 다른 끝 노드를 찾고 그 경로의 길이를 구한다.
k2, distance2 = BFS(k)
print(distance2)