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

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

 

# 백준 1167: 트리의 지름
# 트리의 지름은 가장 먼 두 노드 사이의 거리입니다.

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

V = int(input())  # 노드의 개수

# 노드 연결 정보 저장: nodes[i]에는 (자식 노드 번호, 거리) 쌍이 들어감
nodes = [[] for _ in range(V+1)]
# 각 노드에서 자식 노드를 통해 얻을 수 있는 최대 거리 저장
costs = [[] for _ in range(V+1)]

# 트리 입력 처리
for _ in range(V):
    vals = list(map(int, input().split()))
    parent = vals[0]
    # 입력 형식: 정점 번호, (연결된 정점, 거리)..., -1
    for i in range(1, len(vals)-1, 2):
        child_num, child_cost = vals[i], vals[i+1]
        nodes[parent].append((child_num, child_cost))

# DFS 함수: 현재 노드에서 자식 노드로 이동하며 최대 거리 계산
def dfs(curr_node, prev_node):
    child_cnt = len(nodes[curr_node])  # 연결된 자식 수
    if child_cnt:
        for i in range(child_cnt):
            child = nodes[curr_node][i]
            child_num = child[0]
            child_cost = child[1]

            if prev_node != child_num:
                # 부모 노드로 되돌아가지 않도록 방지 (양방향 그래프이기 때문)
                dfs(child_num, curr_node)
                # 자식 노드로부터 얻은 최대 거리 + 현재 간선 비용 저장
                costs[curr_node].append(child_cost + max(costs[child_num]))
            else:
                # 부모로 가는 경로는 제외 (0으로 처리)
                costs[curr_node].append(0)
    else:
        # 리프 노드일 경우 0 저장
        costs[curr_node].append(0)

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

dfs(1, 0)  # 루트 노드를 1로 두고 탐색 시작

# 각 노드에서 최대 거리 2개를 뽑아 합한 값으로 지름 계산
for i in range(1, V+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(ans)  # 결과 출력

 

 

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

 

위 문제와 동일한 문제이다. 사실 트리의 지름은 어느 노드에서 시작하든 구할 수 있다는게 핵심!