다익스트라 알고리즘 (Python)

개념

다익스트라 알고리즘은 최단 거리를 구할 때 유용한 알고리즘입니다. 특히 한 지점에서 다른 지점으로 이동할 때 최단거리를 구할 때 유용한 알고리즘입니다. 먼저 그림을 보겠습니다.

나동빈님의 이것이 취업을 위한 코딩테스트이다 다익스트라 알고리즘 그래프를 참조했습니다.

위와 같은 그래프가 있을때 우선 distance 테이블을 생성합니다. 그 후 모든 거리를 INF 즉 무한으로 초기화합니다.

노드번호	1	2	3	4	5	6
거리	INF	INF	INF	INF	INF	INF

그 후 start 지점을 1로 설정할 경우를 예를 들어서 설명하겠습니다. 1에서 1로 가는 경우 거리는 0이기 때문에 0으로 설정합니다. 그 후 1번 노드에서 갈 수 있는 지점인 2, 3, 4번 노드로 가는 거리가 각각 2, 5, 1 이고 이들은 모두 INF보다 작으므로 INF를 2, 5, 1로 바꿔줍니다. 그러면 다음과 같습니다. 

 

노드번호	1	2	3	4	5	6
거리	0	2	5	1	INF	INF

그 후 거리가 가장 작은 노드를 선택합니다. 위의 표에서 거리가 작은 노드는 거리가 1인 노드 4 이므로 4를 시작 지점으로 한 후 4에서 연결된 노드를 찾습니다. 4에서 갈 수 있는 노드는 3, 5가 있고 각각의 거리는 3, 1입니다. 1에서 3까지의 거리는 위의 표에서 5라고 나와있습니다. 그러나 1에서 4를 거쳐서 5를 가는 거리는 (1 + 3)이므로 4입니다. 따라서 4로 바꿔주고 5번으로 가는 거리는 원래 INF에서 (1 + 1) 즉 2가 더 작으므로 2로 교체합니다. 

 

노드번호	1	2	3	4	5	6
거리	0	2	4	1	2	INF

이 과정을 전체 노드를 다 반복을 하게 되면 start인 1번 노드에서 모든 노드로 가는 최소 거리를 알 수 있습니다. 소스코드는 다음과 같습니다. 여기서 갈 수 있는 노드 중에 거리가 가장 작은 노드를 고르는 방법은 heapq를 이용해서 구현합니다. heapq를 이용하면 최악의 경우에도 O(ElogV)로 동작하기 때문에 heapq를 이용합니다.

---

소스코드

 

import heapq
import sys
input = sys.stdin.readline
INF = int(1e9)
 
n, m = map(int, input().split())
start = int(input())
graph = [[] for _ in range(1 + n)]
distance = [INF] * (1 + n)
 
for _ in range(m):
  a, b, c = map(int, input().split())
  graph[a].append((b, c))
 
def dijkstra(start):
  q = []
  heapq.heappush(q, (0, start))
  distance[start] = 0
  while q:
    dist, now = heapq.heappop(q)
    if distance[now] < dist:
      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, 1 + n):
  if distance[i] == INF:
    print("INF")
  else:
    print(distance[i])