[백준 1167번] 트리의 지름 (Python/파이썬)

https://www.acmicpc.net/problem/1167

1167번: 트리의 지름

트리의 지름: 노드와 노드 사이의 거리가 가장 긴 거리

이 문제에서의 트리는 가중치가 있는 트리이므로 거리를 구할때 그냥 +1씩 하는 것이 아니라 가중치를 저장했다가 가중치를 더해주어야한다.

이 문제에서 중요한 개념이 하나 있다

"임의의 노드에서 가장 긴 경로로 연결되어 있는 노드는 트리의 지름에 해당하는 두 노드중 하나이다."

위의 개념을 이용하여 단 2번의 탐색으로도 가장 긴 거리 즉 지름을 구할 수 있다.

처음에 아무 노드나 하나 BFS를 돌린다. BFS의 결과로 가장 먼 거리(a)를 가진 노드가 나올 것이고 그 노드에 대해서 또 BFS를 돌린다.

그럼 노드에 대한 가장 먼 거리(b)가 나올 것이다.

답은 a와 b 중에서 큰 값이 가장 긴 거리 즉 지름이 된다.

코드는 다음과 같다.

 

import sys
from collections import deque
input = sys.stdin.readline
n = int(input())
tree = [[] for _ in range(n+1)]
diameter = 0
distance = [0] * (n+1)
visited = [False] * (n+1)
for _ in range(n):
  info = list(map(int,input().split()))
  idx = 0
  node = info[idx]
  idx += 1
  while info[idx] != -1:
    tmp_node = info[idx]; tmp_weight = info[idx+1]
    tree[node].append((tmp_node,tmp_weight))
    idx += 2

def BFS(v):
  queue = deque([])
  queue.append(v)
  visited[v] = True
  while queue:
    node = queue.popleft()
    for i in tree[node]:
      n,w = i
      if not visited[n]:
        visited[n] = True
        distance[n] = distance[node] + w
        queue.append(n)
BFS(1)

for i in range(1,n+1):
  if diameter < distance[i]:
    diameter = distance[i]
    longest = i
distance = [0] * (n+1)
visited = [False] * (n+1)
BFS(longest)
tmp = max(distance)
if diameter < tmp:
  diameter = tmp
print(diameter)

 

 

 

 

 

문제의 구현이 어렵다기 보다는 위에 언급했던 개념을 알면 단순 BFS이므로 쉽게 풀어낼 수 있는 문제이다.