[ 백준 / 파이썬 ] 골드 2 - 1167. 트리의 지름

난이도

문제

트리의 지름이란, 트리에서 임의의 두 점 사이의 거리 중 가장 긴 것을 말한다. 트리의 지름을 구하는 프로그램을 작성하시오.

입력

트리가 입력으로 주어진다. 먼저 첫 번째 줄에서는 트리의 정점의 개수 V가 주어지고 (2 ≤ V ≤ 100,000)둘째 줄부터 V개의 줄에 걸쳐 간선의 정보가 다음과 같이 주어진다. 정점 번호는 1부터 V까지 매겨져 있다.

먼저 정점 번호가 주어지고, 이어서 연결된 간선의 정보를 의미하는 정수가 두 개씩 주어지는데, 하나는 정점번호, 다른 하나는 그 정점까지의 거리이다. 예를 들어 네 번째 줄의 경우 정점 3은 정점 1과 거리가 2인 간선으로 연결되어 있고, 정점 4와는 거리가 3인 간선으로 연결되어 있는 것을 보여준다. 각 줄의 마지막에는 -1이 입력으로 주어진다. 주어지는 거리는 모두 10,000 이하의 자연수이다.

출력

첫째 줄에 트리의 지름을 출력한다.

예제

입력	출력
5 1 3 2 -1 2 4 4 -1 3 1 2 4 3 -1 4 2 4 3 3 5 6 -1 5 4 6 -1	11

풀이

그냥 일반적인 DFS, BFS 로는 풀 수 없다. 시간초과가 발생하기 때문..

그래서, 유튜브 영상을 참고하여 이런 문제 유형을 푸는 방법을 알아냈다.

트리의 지름으로 구성될 두 노드는, 반드시 가장 긴 간선을 가지고 있다.

이러한 아이디어를 가지고 해당 문제를 풀어야 한다.

최대 거리가 트리의 지름이 될 것이기 때문에, 지름을 구성하게 될 간선에는 반드시 가장 긴 간선을 가지고 있을 수 밖에 없다.

출발점에서 가장 긴 간선을 가지는 노드를 찾고,
해당 노드에서의 최장 거리가 트리의 지름이 된다.

코드

from sys import stdin
from collections import deque

input = stdin.readline

def bfs(q, max_node, max_cost, tree, visited):
    while q:
        node, cost = q.popleft()
        if max_cost < cost:
            max_node = node
            max_cost = cost

        for new_node, new_cost in tree[node]:
            if not visited[new_node]:
                visited[new_node] = True
                q.append((new_node, cost + new_cost))

    return max_node, max_cost

def solution():
    V = int(input())
    tree = [[] for _ in range(V + 1)]
    for _ in range(V):
        arr = list(map(int, input().split()))
        S = arr[0]
        arr = arr[1:-1]
        for i in range(0, len(arr) - 1, 2):
            tree[S].append((arr[i], arr[i + 1]))

    max_node = 0
    max_cost = 0
    visited = [False] * (V+1)
    visited[1] = True
    q = deque([(1, 0)])

    max_node, max_cost = bfs(q, max_node, max_cost, tree, visited)

    q = deque([(max_node, 0)])
    visited = [False] * (V + 1)
    visited[max_node] = True

    max_node, max_cost = bfs(q, max_node, max_cost, tree, visited)


    print(max_cost)


solution()