본문 바로가기
Algorithm/Graph

[파이썬] 백준 1967 : 트리의 지름 (골드4)

by 베짱이28호 2024. 8. 1.

[파이썬] 백준 1967 : 트리의 지름 (골드4)


풀이

방향성 생각

  • 모르면 노드마다 $N^2$로 돌리겠지만, BFS 두 번으로 해결이 가능하다.
  • 루트노드인 1번 노드에서 BFS를 돌리고, 가장 먼 지점에서 다시 한 번 BFS


전체코드

from collections import deque
import sys
input = lambda : sys.stdin.readline().rstrip()

N = int(input())
G = [[] for _ in range(N+1)]
for _ in range(N-1):
    a,b,cost = map(int,input().split())
    G[a].append((b,cost))
    G[b].append((a,cost))

def bfs(x):

    global N

    Q = deque([(x,0)])
    V = [-1]*(N+1)
    V[x] = 0
    while Q:
        x,t = Q.popleft()
        for nx,nt in G[x]:
            if V[nx] == -1:
                V[nx] = t+nt
                Q.append((nx,t+nt))
    return V

max_dist = 0
start = 1
for idx,dist in enumerate(bfs(1)):
    if dist > max_dist:
        max_dist = dist
        start = idx

for idx,dist in enumerate(bfs(start)):
    max_dist = max(max_dist,dist)
print(max_dist)

코멘트

  • 98%에서 Node 개수가 1개인 케이스가 있는듯?
    • 마지막 start = 1을 할당 안해주고, for문 내부에서만 활용하면 name error가 있다

댓글