[백준알고리즘] 1967번: 트리의 지름 -Python

[백준알고리즘] 1967번: 트리의 지름 -Python

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

1967번: 트리의 지름

이전에 풀었던 "1167번 트리의 지름" 문제와 거의 유사하다. 오히려 이 문제가 더 쉽다.

 

트리의 지름을 구하는 방식은 1167번 링크에서 확인을 하면 된다.

 

이 문제가 1167번 보다 쉬운 이유는 입력에 순서가 정해져 있기 때문이다. 부모가 작은 순서대로, 자식이 작은 순서대로 입력이 주어지기 때문에 트리는 min heap tree와 유사한 형태가 된다. 게다가 자식 또한 작은 값이 무조건 왼쪽에 위치하는 정렬된 형태를 갖게 된다.

 

그리고 문제에는 언급이 안되어있는데 정답처리된 코드를 살펴보니 정점도 그냥 n개의 개수일 뿐만 아니라 1~n까지의 정점만 주어지나 보다.

 

입력의 순서가 정해져 있다 보니 풀었던 것처럼 딕셔너리를 사용할 필요가 없고 더 간단히도 가능했겠지만 이전 문제와 유사하게 풀며 익히기 위해서 아래처럼 풀었다. 물론 속도는 다른 코드들에 비해 더 나왔다.

import sys
input = sys.stdin.readline
 
def farthest(i):
    stack = tree[i][:]
    visit = [False] * 10001
    visit[i] = True
    node, dist = 0, 0
 
    for sn, sd in stack:
        visit[sn] = True
        if sd > dist:
            node = sn
            dist = sd
 
    while stack:
        sn, sd = stack.pop()
        link = tree[sn]
        for ln, ld in link:
            if not visit[ln]:
                visit[ln] = True
                now = sd + ld
                stack.append((ln, now))
                if now > dist:
                    node = ln
                    dist = now
    
    return node, dist
 
n = int(input())
tree = {}
key = []
for _ in range(n-1):
    v1, v2, weight = map(int, input().split())
    if not v1 in key:
        tree[v1] = []
        key.append(v1)
    if not v2 in key:
        tree[v2] = []
        key.append(v2)
    tree[v1].append((v2, weight))
    tree[v2].append((v1, weight))
 
edge, dist = farthest(1)
print(farthest(edge)[1])

 

잘못된 점이나 부족한 점 지적해주시면 감사하겠습니다