[Python] 백준 1167번: 트리의 지름

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

1167번: 트리의 지름

import sys
input = sys.stdin.readline


def setDistances(node):
    global distances

    for n, d in tree[node]:
        if distances[n] == -1:
            distances[n] = distances[node] + d

            setDistances(n)


if __name__ == '__main__':
    V = int(input())

    # 입력 기반으로 tree 구성
    tree = [[] for i in range(V+1)]

    for _ in range(V):
        curNums = list(map(int, input().split()))
        curNumsLen = len(curNums)

        for i in range(1, curNumsLen-1, 2):
            tree[curNums[0]].append((curNums[i], curNums[i+1]))

    # 단말 노드 아무거나 지정
    for i in range(1, V+1):
        if len(tree[i]) == 1:
            terminalNode = i
            break

    # 단말 노드를 기준으로 탐색하여, 거리가 가장 먼 노드의 정보를 얻음
    distances = [-1]*(V+1)
    distances[terminalNode] = 0
    setDistances(terminalNode)
    maxDistanceNode = distances.index(max(distances))

    # 가장 먼 거리에 해당하는 노드에서 탐색하여 최대 길이를 구함
    distances = [-1]*(V+1)
    distances[maxDistanceNode] = 0
    setDistances(maxDistanceNode)
    maxDistance = max(distances)

    print(maxDistance)

우선 단말 노드에서 시작해서 dfs나 bfs를 돌리면서 누적 거리를 저장하고 가장 큰 거리를 출력해주면 된다.

하지만 각 단말 노드마다 최대 거리가 다르다. 예제 입력 1을 보면, 만약 탐색을 2에서 시작하면 최대 거리는 10이 되지만, 탐색을 1이나 5에서 시작하면 11이 된다.

우선 모든 단말 노드에 대해 탐색을 수행하는 것은 그리 좋지 않다. Pypy3으로 돌려도 시간 초과가 나기 때문이다.

이에 대한 해결책은 간단하다.

1. 하나의 단말 노드에서 탐색 시도

2. 탐색하여 얻은 정보에서 가장 먼 거리에 있는 노드를 찾음

3. 해당 노드를 기준으로 탐색한 뒤, 최대 거리 출력

이것을 나는 dfs로 구현하였다. 유의할 점은 탐색을 할 때마다 distances를 초기화시켜줘야 한다는 점이다.