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

문제

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

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

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

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

 

기록

(트리 3일차) 노드별로 지름 최댓값 저장하는 방법을 잘못 구현했는데 고치니까 맞았다.

 

코드

from sys import stdin, setrecursionlimit
from collections import defaultdict
setrecursionlimit(10 ** 9)

# 경로 최댓값 찾으면서 노드 별 지름 저장
def find_max(tree, start, diadict):
    # base case
    if start not in tree : return 0
    # step
    rootmax = 0
    rootmax2 = 0
    for node, w in tree[start].items():
        nodemax = tree[start][node] + find_max(tree, node, diadict)
        if nodemax>rootmax :
            rootmax2 = rootmax
            rootmax = nodemax
        elif nodemax>rootmax2 :
            rootmax2 = nodemax
    diadict[start] = rootmax + rootmax2
    return rootmax

# 변수 초기화
n = int(stdin.readline())
tree = defaultdict(dict)
for i in range(n-1):
    u, v, w = list(map(int,stdin.readline().split()))
    tree[u][v] = w

# 지름 최댓값찾기
diadict = {}
find_max(tree, 1, diadict)
diamax = 0
for i in diadict.values():
    if diamax<i : diamax=i
print(diamax)

1. 입력값은 {루트 : {노드 : weight, 노드 : weight ...}} 형식으로 저장

2. 루트노드부터 최하위노드로 가며 가지들 중 최댓값을 반환하는 재귀함수를 수행하며, 따로 현재노드를 중심으로 하는 지름(최댓값 2개 가지의 합)을 dict에 저장하기

3. 지름들을 저장한 dict에서 가장 큰 지름 반환하기