최단경로 (다익스트라)

백준 1753번: 최단경로

아이디어

다익스트라 알고리즘 문제다. 시작 노드를 선택하고 그 노드부터 연결된 노드를 탐색하며 최단 거리를 저장해 둔 리스트를 업데이트한다. 연결된 노드중에 비용이 가장 적은 노드를 다음으로 우선으로 탐색한다(그리디 알고리즘 느낌). 비용이 가장 적은 노드를 우선으로 선택하기 위해, 우선순위 큐 (최소 힙)를 사용한다. 우선순위 큐를 사용하지 않고 비용이 가장 적은 노드를 탐색하면 O(N), 우선순위 큐를 사용해서 비용이 가장 적은 노드를 탐색하면 O(logN)이다.

코드

import heapq
import sys

input = sys.stdin.readline

INF = 987654321

V, E = map(int, input().split())
start = int(input())

dist = [INF] * (V + 1)
dist[start] = 0

graph = [[] for _ in range(V + 1)]

for _ in range(E):
    u, v, w = map(int, input().split())
    graph[u].append((v, w))

def dijkstra(start):
    q = []
    # "비용"을 기준으로 삽입, 삭제 연산을 해야 한다. 
    # 따라서 튜플의 첫 번째 원소는 비용.
    heapq.heappush(q, (0, start))
    while q:
        c, t = heapq.heappop(q)
        if c > dist[t]:
            continue
        for i in graph[t]:
            if dist[i[0]] > c + i[1]:
                dist[i[0]] = c + i[1]
                heapq.heappush(q, (c + i[1], i[0]))


dijkstra(start)
for n in range(1, V + 1):
    if dist[n] == INF:
        print("INF")
    else:
        print(dist[n])

여담

두 번 정도 헤매다가 맞았다.
일단 우선순위 큐를 사용해야 하는데 그냥 큐를 사용해서 풀었다. 거기서 계속 시간초과가 발생했고 또 우선순위 큐의 삽입, 삭제 연산 기준을 "비용"으로 해야하는데 반대로 노드 번호를 써버려서 계속 시간초과가 발생했다.

우선순위 큐 + "비용" 기준 정렬 기억하도록 하자.