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

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

1967번: 트리의 지름

문제

트리(tree)는 사이클이 없는 무방향 그래프이다. 트리에서는 어떤 두 노드를 선택해도 둘 사이에 경로가 항상 하나만 존재하게 된다. 트리에서 어떤 두 노드를 선택해서 양쪽으로 쫙 당길 때, 가장 길게 늘어나는 경우가 있을 것이다. 이럴 때 트리의 모든 노드들은 이 두 노드를 지름의 끝 점으로 하는 원 안에 들어가게 된다.

이런 두 노드 사이의 경로의 길이를 트리의 지름이라고 한다. 정확히 정의하자면 트리에 존재하는 모든 경로들 중에서 가장 긴 것의 길이를 말한다.

입력으로 루트가 있는 트리를 가중치가 있는 간선들로 줄 때, 트리의 지름을 구해서 출력하는 프로그램을 작성하시오. 아래와 같은 트리가 주어진다면 트리의 지름은 45가 된다.

트리의 노드는 1부터 n까지 번호가 매겨져 있다.

코드

import sys
sys.setrecursionlimit(100000)
n = int(sys.stdin.readline())
tree = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    a, b, c = map(int, sys.stdin.readline().split())
    tree[a].append((b, c))
max_n = 0
def dfs(n, d):
    global max_n
    left, right = 0, 0
    for node, w in tree[n]:
        temp = dfs(node, w)
        if left <= right:
            left = max(left, temp)
        else:
            right = max(right, temp)
    max_n = max(max_n, left + right)
    return max(left + d, right + d)
dfs(1, 0)
print(max_n)

풀이

왼쪽에서 가장 긴 경로와 오른쪽에서 가장 긴 경로를 더한 것이 트리의 지름이 된다. 각각 가장 긴 것을 찾아 max_n 변수에 저장해 출력하면 된다.