다익스트라 알고리즘 문제다. 시작 노드를 선택하고 그 노드부터 연결된 노드를 탐색하며 최단 거리를 저장해 둔 리스트를 업데이트한다. 연결된 노드중에 비용이 가장 적은 노드를 다음으로 우선으로 탐색한다(그리디 알고리즘 느낌). 비용이 가장 적은 노드를 우선으로 선택하기 위해, 우선순위 큐 (최소 힙)를 사용한다. 우선순위 큐를 사용하지 않고 비용이 가장 적은 노드를 탐색하면 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])
두 번 정도 헤매다가 맞았다.
일단 우선순위 큐를 사용해야 하는데 그냥 큐를 사용해서 풀었다. 거기서 계속 시간초과가 발생했고 또 우선순위 큐의 삽입, 삭제 연산 기준을 "비용"으로 해야하는데 반대로 노드 번호를 써버려서 계속 시간초과가 발생했다.
우선순위 큐 + "비용" 기준 정렬 기억하도록 하자.