1 minute read

문제

joowan1108

접근 방법

우선 모든 노드에 대해 dijkstra나 bfs와 같은 거리 알고리즘을 적용하면 시간 초과가 날 것임을 볼 수 있다. 직관적으로 생각했을 때 트리의 지름을 구하는 것이므로 서로 제일 먼 두 끝 점을 구해야 한다. 그러기 위해서는 우선 루트 노드에 대해 dijkstra를 적용하여 root 노드에서 가장 먼 노드를 구했다.

여기서 root node에 대해 dijkstra를 적용하여 제일 먼 두 노드의 합으로 지름을 구할 수 있을 줄 알았는데 더 생각해보면 두 노드가 갈라져있지 않을 수도 있는 case가 존재한다는 것을 알 수 있다…

이어서 root 노드에서 가장 먼 노드로부터 가장 먼 노드까지의 거리를 구하면 된다. 이것이 바로 양 끝점 사이의 거리로 볼 수 있기 때문이다.

정답 코드


import sys
import heapq
input = sys.stdin.readline



"""
root 노드에서 모든 거리를 잰다
root 노드에서 가장 먼 노드는 지름의 한 끝이라고 볼 수 있다

그 노드에서 가장 먼 노드까지의 길이는 어떻게 보면 지름이라고 볼 수 있음
"""


n = int(input())
if n==1:
    print(0)
else:
    graph = [[] for _ in range(n+1)]
    for _ in range(n-1):
        a,b,c = map(int, input().split())
        graph[a].append((b,c))
        graph[b].append((a,c))


    def dijkstra(start):
        q = []
        distances = [sys.maxsize]*(n+1)
        visited = set()
        distances[start] = 0
        heapq.heappush(q, (0,start))
        while q:
            mid_cost, mid_node = heapq.heappop(q)
            if mid_node in visited:
                continue
            visited.add(mid_node)
            for end_node, cost in graph[mid_node]:
                new_distance = distances[mid_node] + cost
                if new_distance < distances[end_node]:
                    distances[end_node] = new_distance
                    heapq.heappush(q,(new_distance, end_node))
        return distances


    #가장 먼 노드 찾기
    distances = dijkstra(1)
    distances[0] = 0
    max_d = 0
    max_node = 0
    for i in range(n+1):
        if max_d < distances[i]:
            max_d = distances[i]
            max_node = i


    distances = dijkstra(max_node)
    distances[0] = 0

    print(max(distances))

느낀 점

edge case에 대해 충분히 고민하는 습관이 필요한 것 같다.

You May Also Enjoy