백준 1967 트리의 지름 (파이썬)

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

1967번: 트리의 지름

트리의 끝 지점만을 저장한 후 bfs로 끝 지점 간의 길이를 구한 후 최대값을 구하였더니 시간초과가 떴다

 

1. root로부터의 모든 경로 중에서 최대값을 구한다

2. 최대값에 해당하는 경로의 끝 지점 인덱스를 구한다

3. 해당 인덱스로부터 이동 경로의 최대값을 구하면 트리의 최대 지름을 구할 수 있다.

 

root로부터의 최대값을 구할 수 있는 경로와 트리의 최대 지름은 겹치는 경로가 있기 때문에 위와 같은 풀이가 가능하다

from collections import deque
import sys

input = sys.stdin.readline

def bfs(x, mode):
    q = deque()
    q.append(x)
    c = [-1 for _ in range(n)]
    c[x] = 0
    while q:
        x = q.popleft()
        for w, nx in a[x]:
            if c[nx] == -1:
                c[nx] = c[x] + w
                q.append(nx)
    if mode == 1:
        return c.index(max(c))
    else:
        return max(c)

n = int(input())
a = [[] for _ in range(n)]

for i in range(n-1):
    x, y, w = map(int, input().split())
    a[x-1].append([w, y-1])
    a[y-1].append([w, x-1])
print(bfs(bfs(0, 1), 2))