백준 1967 [트리의 지름] 파이썬

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

# 백준 1967: 트리의 지름 문제 풀이

import sys
sys.setrecursionlimit(10**7)  # 재귀 깊이 제한 증가
input = sys.stdin.readline

# 노드 수 입력
N = int(input())

# 각 노드의 자식 노드 및 간선 정보를 저장할 리스트
nodes = [[] for _ in range(N+1)]   # 인접 리스트 (노드 간 연결 정보)
costs = [[] for _ in range(N+1)]   # 각 노드에서 자식으로 가는 거리 누적값 저장용

# 트리 간선 정보 입력
for i in range(N-1):
    parent, child, val = map(int, input().split())
    nodes[parent].append((child, val))  # 부모 → 자식 방향 저장

# DFS로 각 노드에서의 최대 거리 계산
def dfs(node):
    child_cnt = len(nodes[node])  # 자식 노드 수 확인
    if child_cnt:
        # 자식 노드가 있다면
        for i in range(child_cnt):
            child = nodes[node][i]
            child_num = child[0]
            child_cost = child[1]
            dfs(child_num)  # 자식 노드 먼저 탐색
            # 자식으로부터의 거리 + 자식 노드의 최대 거리 저장
            costs[node].append(child_cost + max(costs[child_num]))
    else:
        # 리프 노드일 경우 0 저장
        costs[node].append(0)

ans = 0  # 트리의 지름(최대 거리)

dfs(1)  # 루트 노드부터 탐색 시작

# 각 노드에서 가장 큰 거리 2개를 더한 값으로 지름 계산
for i in range(1, N+1):
    b1 = 0  # 가장 큰 거리
    b2 = 0  # 두 번째로 큰 거리
    for cost in costs[i]:
        if cost > b2:
            if cost > b1:
                b2 = b1
                b1 = cost
            else:
                b2 = cost
    ans = max(b1 + b2, ans)  # 가장 긴 두 경로의 합으로 지름 갱신
    # print(costs[i])  # 디버깅용 출력

# 결과 출력
print(ans)

자식의 노드가 여러개인 경우를 생각해야한다.

자식 노드 두개의 합이 최대가 되는 노드를 찾으면 정답이다.