백준 1167 파이썬 - 트리의 지름

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

1167번: 트리의 지름

문제

트리의 지름이란, 트리에서 임의의 두 점 사이의 거리 중 가장 긴 것을 말한다. 트리의 지름을 구하는 프로그램을 작성하시오.

코드

import sys
from collections import deque
v = int(sys.stdin.readline())
tree = [[] for _ in range(v + 1)]
for _ in range(v):
    arr = list(map(int, sys.stdin.readline().split()))
    for i in range(1, len(arr) - 1, 2):
        tree[arr[0]].append([arr[i], arr[i + 1]])
def bfs(i):
    q = deque()
    visited = [-1] * (v + 1)
    q.append(i)
    visited[i] = 0
    max_n = [0, 0]
    while q:
        n = q.popleft()
        for a in tree[n]:
            if visited[a[0]] == -1:
                visited[a[0]] = visited[n] + a[1]
                q.append(a[0])
                if max_n[0] < visited[a[0]]:
                    max_n[0] = visited[a[0]]
                    max_n[1] = a[0]
    return max_n
array = bfs(1)
print(bfs(array[1])[0])

풀이

1에서 부터 가장 멀리있는 노드를 구하고 그 후 그 노드를 한번 더 bfs로 돌려 가장 멀리있는 노드를 구하면 된다.