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

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

1967번: 트리의 지름

 

import sys
input = sys.stdin.readline


def getFarthestNodeAndPrefixWeight(start, n):
    todo = [(start, 0)]
    isVisited = [False]*(n+1)

    farthestNode = start
    maxWeight = 0

    while todo:
        node, weight = todo.pop(0)
        isVisited[node] = True

        if maxWeight < weight:
            maxWeight = weight
            farthestNode = node

        for n, w in nodes[node]:
            if isVisited[n] == False:
                todo.append((n, weight + w))

    return (farthestNode, maxWeight)


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

    nodes = {i: [] for i in range(1, n+1)}

    for _ in range(n-1):
        a, b, c = map(int, input().split())
        nodes[a].append((b, c))
        nodes[b].append((a, c))

    farthestNodeFromRoot, _ = getFarthestNodeAndPrefixWeight(1, n)
    _, diameter = getFarthestNodeAndPrefixWeight(farthestNodeFromRoot, n)
    print(diameter)

 

우선 입력을 그래프를 처리하듯이, dictionary로 받아준다. 이후에는 아래와 같은 순서로 진행하면 된다.

1. 루트(1)로부터 거리가 가장 먼 노드(farthestNodeFromRoot)를 찾는다.

2. 찾아낸 노드로부터 거리가 가장 먼 노드를 찾고, 이 때 탐색하면서 누적된 weight를 출력한다.

 

매우 단순한 구조이며, 이를 위한 함수인 getFarthestNodeAndPrefixWeight(start, n)을 설계한다. 해당 함수는, 간단한 bfs에서, maxWeight을 통해 가장 거리가 먼 노드(farthestNode)와 해당 노드까지의 거리(maxWeight)를 갱신시켜주는 부분이 추가된 형태이다. bfs 이후에 갱신된 farthestNode와 maxWeight를 리턴해주면 된다.