[백준 1967] 트리의 지름(Python)

1. 요약

간선에 가중치가 있는 트리에서 임의의 두 노드를 선택할 때, 두 노드 사이의 경로가 가장 긴 것을 찾는 문제

2. 아이디어

긴 경로를 갖는 후보는 두 노드가 모두 리프노드 일 때다.(리프노드가 1개밖에 없으면 루트노드와 리프노드 사이의 경로가 가장 길다.) 이들 중 가장 긴 경로를 찾으면 된다.모든 트리의 노드가 갖는 서브트리에 대해 dfs를 한다.

dfs의 결과는 서브트리의 루트노드와 리프노드사이의 가장 긴 경로이다. 만약 어떤 노드가 여러개의 서브트리를 갖는다면, 각 결과를 임시배열에 저장하고 가장 큰 값 2개를 가져와서 더하면 된다.

3. 코드

import sys
from collections import defaultdict

sys.setrecursionlimit(int(1e4))
n = int(sys.stdin.readline().rstrip())
tree = defaultdict(list)
answer = 0

for _ in range(n-1):
    a, b, w = map(int, sys.stdin.readline().rsplit())
    tree[a].append((b, w))

def solve(node, cost):
    if not tree[node]: return cost
    total = 0
    m = 0
    for next_node, w in tree[node]:
        m = solve(next_node, cost+w)
        if total < m: total = m
    
    return total

for node in range(n):
    tmp = []
    for next_node, w in tree[node]:
        tmp.append(solve(next_node, w))
    
    tmp.sort(reverse=True)
    total = sum(tmp[:2])
    if answer < total: answer = total

print(answer)