[백준/파이썬] 1967번: 트리의 지름 핵심 정리

문제 보기 : https://www.acmicpc.net/problem/1967

 

트리의 지름 문제의 핵심 아이디어

사이클이 없는 무방향 그래프인 트리가 존재하고, 이 때 어떤 두 노드를 선택해도 둘 사이의 경로는 오직 하나만 존재하는데요.

 

이 때 임의의 두 노드를 선택해 양쪽으로 당길 때 가장 길게 늘어나는 경우, 이 두 노드 사이 경로의 길이를 '트리의 지름'이라고 합니다. 즉 임의의 두 노드 간의 가장 긴 경로 합을 구하는 것이 이 문제의 요구사항입니다.

 

그러고 나서 문제는 트리의 지름을 이루는 두 노드와 그 안에 있는 나머지 노드들로 이루어진 원 형태의 그림을 보여주는데요, 이 그림이 바로 문제를 푸는 핵심이라고 할 수 있습니다.

 

 

그림을 보면 트리의 모든 노드는 이 원 안에 위치합니다. 다시말해 어떤 노드를 기준으로 하던 간에, 해당 노드부터 가장 긴 경로를 가진 노드는 트리의 지름을 이루는 두 노드 중 하나가 된다는 것이죠.

 

이렇게 트리의 지름을 이루는 두 노드 중 하나를 알게 된다면, 다시 방금 구한 노드를 기준으로 해당 노드부터 가장 긴 경로를 가진 노드를 구함으로써 트리의 지름을 이루는 나머지 하나의 노드와 그 경로 합(트리의 지름)을 알 수 있게 됩니다.

 

 

예시 분석을 통해 이해하기

문제에서는 임의의 트리에 대한 예시를 함께 들고 있는데요. 정말 아무 노드나 선택해서 위의 핵심이 유효한지 확인해보겠습니다.

 

위의 트리에서 8번 노드를 기준으로 dfs 혹은 bfs를 통해 가장 긴 경로 합을 이루는 노드를 구한다고 하겠습니다.

그러면 위와 같이 A, B, C, D, E의 경로가 존재할 것입니다. 직접 계산해보면 이 중 경로 B(9번 노드)가 가장 긴 경로 합을 가지고 있음을 알 수 있는데요.

 

이 때 9번 노드로부터 가장 긴 경로 합을 다시 구한다고 하면, 단순히 도로 경로 B를 타고 8번 노드가 나오는 것이 아닌가? 할 수 있습니다. 하지만 이 때 시작 노드에 따라 가장 긴 경로 합을 가지고 있을 후보들이 달라진다는 점을 고려해야 합니다.

위의 경로들은 9번 노드로부터 가장 긴 경로 합을 가지고 있는 경로들입니다. 기존의 8번 노드로 가는 경로(경로 D)도 후보군에 포함되어 있습니다만, 가장 긴 경로 합을 가지는 것은 12번 노드로 가는 경로(경로 C)입니다. 그리고 이 9번 노드와 12번 노드의 경로 합이 저희들이 원하는 트리 지름이 됩니다. 

---

예시를 하나 더 들어볼까요? 경로에 대한 그림은 생략하고 이번에는 5번 노드를 선택해 가장 긴 경로 합을 이루는 노드를 구해보도록 하겠습니다. 후보군은 (5-9), (5-10), (5-3-6-11), (5-3-6-12), (5-3-1-2-4-7), (5-3-1-2-4-8) 인데 이중 가장 긴 경로 합은 (5-3-6-12)로 12번 노드가 가장 긴 경로 합을 이루는 노드입니다.

 

다시 이 12번 노드를 선택해 후보군을 추려보면 (12-6-11), (12-6-3-5-9), (12-6-3-5-10), (12-6-3-1-2-4-7), (12-6-3-1-2-4-8)이 있고, 이 중 가장 긴 경로 합을 가지는 것은 (12-6-3-5-9)가 됩니다. 결론적으로 9번, 12번 노드의 경로 합(45)이 트리 지름인 것이죠.

 

 

구현하기

위의 핵심 내용을 바탕으로, bfs 또는 dfs를 통해 임의의 시작 노드로부터 방문하지 않은 노드를 방문해 경로 합을 구해줌으로써 시작 노드에 대한 각 노드의 경로 합을 구할 수 있는데요. 이때 임의의 시작 노드는 어떤 노드를 선택하던 상관이 없습니다.

 

위의 과정을 통해 트리 지름을 이루는 노드 중 하나의 노드를 구할 수 있고, 다시 이를 시작 노드로 하여 bfs(또는 dfs)를 수행하면 트리 지름을 이루는 나머지 노드와 두 노드의 거리(트리 지름)를 구할 수 있습니다. 쉽게 말해 dfs(또는 bfs)를 두 번 수행하면 됩니다.

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

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

def bfs(start):
    q = deque()
    q.append(start)
    
    visited = [-1] * (n+1)
    visited[start] = 0
    
    while q:
        now = q.popleft()
        
        for node, weight in graph[now]:
            if visited[node] == -1:
                visited[node] = visited[now] + weight
                q.append(node)
    # 최대 경로를 이루는 노드, 최대 경로 합 반환
    return visited.index(max(visited)), max(visited)

# node1, _ = bfs(1)
# node2, diameter = bfs(node1)
# print(diameter)
# 위의 세 줄을 축약하면 아래와 같다.
print(bfs(bfs(1)[0])[1])