다익스트라 알고리즘은 그래프에서 최단 거리를 구하는 알고리즘으로 에지는 모두 0이상의 양수여야 한다.
시간복잡도: O(ElogV) V(노드 수) E(에지 수)
import sys
import heapq
input = sys.stdin.readline
INF = sys.maxsize
n,m = map(int,input().split())
start = int(input())
distance = [INF] * (n+1)
graph = [[] for _ in range(n+1)]
for _ in range(m):
s,e,cost = map(int,input().split())
# s번 노드에서 e번 노드로 가는 비용은 cost이다
graph[s].append((e,cost))
def dijkstra(start):
q = []
distance[start] = 0
heapq.heappush(q, (0, start))
while q:
# 가장 최단 거리가 짧은 노드에 대한 정보 꺼내기
dist, now = heapq.heappop(q)
#현재 노드가 이미 처리된 적이 있는 노드라면 무시
if dist > distance[now]:
continue
# 현재 노드와 연결된 다른 인접한 노드들을 확인
for i in graph[now]:
cost = dist + i[1]
# 현재 노드를 거쳐서, 다른 노드로 이동하는 거리가 더 짧은 경우
if cost < distance[i[0]]:
distance[i[0]] = cost
heapq.heappush(q, (cost, i[0]))
dijkstra(start)
for i in range(1, n+1):
if distance[i] == INF:
print("INF")
else:
print(distance[i])
출처
Do it algorithm