[백준 / 2132] 나무 위의 벌레 - Python

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

2132번: 나무 위의 벌레

<풀이>

트리의 지름이란 트리를 양쪽으로 잡아당겼을 때 가장 긴 거리이다.

 

이 트리의 지름을 구하는 방법은 임의의 한점에서 가장 먼 점(A)을 구하고, 해당 점(A)에서 가장 먼 점(B)를 연결하는 길이(AB)가 가장 큰 길이이다.

 

이 특징을 이용해 해당 문제를 풀 수 있다.

 

주의해야 할 점은 길이가 같은 지름이 여러개 존재할 수 있으므로, 길이가 같은 것을 다 구해서 모두 구해줘야 한다.

 

 

<전체 코드>

import sys
from collections import deque

N = int(sys.stdin.readline())
fruits = list(map(int, sys.stdin.readline().split()))
graph = {i: [] for i in range(1, N+1)}
for _ in range(N-1):
    x, y = map(int, input().split())
    graph[x].append(y)
    graph[y].append(x)


def DFS(n):
    distances = [-1]*(N+1)
    que = deque([[n, n]])
    distances[n] = [fruits[n-1], [n, n]]

    maxDis = 0
    nodeList = []
    while que:
        route = que.pop()
        now = route[0]

        for next in graph[now]:
            if distances[next] == -1:
                temp = route[:]
                temp[0] = next
                dis = distances[now][0]+fruits[next-1]
                distances[next] = [dis, temp]
                que.append(temp)
                if maxDis < dis:
                    maxDis = dis
                    nodeList = [next]
                elif maxDis == dis:
                    nodeList.append(next)
    same = {}
    for node in nodeList:
        same[node] = [distances[node][0], min(distances[node][1])]
    return same


dis = DFS(1)
maxDis = -1
ansNode = -1
for node in dis:
    finalDis = DFS(node)
    for ans in finalDis:
        if maxDis < finalDis[ans][0]:
            maxDis = finalDis[ans][0]
            ansNode = finalDis[ans][1]
        if maxDis == finalDis[ans][0]:
            ansNode = min(finalDis[ans][1], ansNode)
if N == 1:
    print(fruits[0], 1)
else:
    print(maxDis, ansNode)