최단 경로 - 다익스트라 알고리즘 with Python (선형 탐색을 이용한 간단한 구현)

다익스트라 알고리즘은 시작 노드로부터 모든 다른 노드로의 최단 경로를 구한다.

다익스트라 알고리즘 동작 원리
1. 출발 노드 설정
2. 최단 거리 테이블 초기화
3. 방문하지 않은 노드 중 최단 거리가 가장 짧은 노드를 선택
4. 해당 노드를 거쳐 다른 노드로 가는 비용을 계산하여 최단 거리 테이블 갱신
5. 모든 노드를 방문할 때까지 3번, 4번을 반복

최초에 최단 거리 테이블을 999,999,999 (약 10억)으로 초기화하는데,
자리수를 기억하기 힘들다면 987,654,321로 기억하면 된다.
또한 간편하게 1e9를 사용할 수도 있으나 파이썬에서 이를 정수가 아닌 실수로 인식하기 때문에 int(1e9)를 사용한다.

다익스트라 알고리즘을 구현하는 방법은 2가지이다.

• 구현이 쉽지만 동작이 느린 방법
• 구현이 까다롭지만 동작이 비교적 빠른 방법

이번 포스트에서는 첫 번째 방법을 다루기로 하고, 다음 포스트에서 priority queue를 사용한 두 번째 방법을 다루기로 한다.

아래 코드는 구현이 쉽지만 느린 방법으로, O(V2)의 시간 복잡도를 가진다.

import sys
input = sys.stdin.readline
INF = int(1e9)

#n, m = 노드 개수, 간선 개수
n, m = map(int, input().split())
start = int(input()) #시작 노드
# 각 노드에 대해 (연결 노드, 거리) 값을 담기 위한 그래프
graph = [ [] for _ in range(n+1) ]
#다익스트라 알고리즘 수행 시 방문 여부 체크를 위한 리스트
visited = [0] * (n + 1) 
#최단 거리 테이블 
distance = [INF] * (n + 1)

for _ in range(m):
  #a노드에서 b노드로 가는 간선의 거리 c
  a, b, c = map(int, input().split())
  graph[a].append((b, c))

#최단 거리 테이블의 최소값을 얻기 위한 함수
def get_smallest_node():
  idx = 0 
  min_value = INF
  for i in range(1, n + 1):
    if distance[i] < min_value and visited[i] == 0:
      min_value = distance[i]
      idx = i
  return idx

def dijkstra(start):
  #시작 노드에서 시작 노드까지의 거리 0 저장
  distance[start] = 0
  #시작 노드를 포함한 n개의 노드 방문
  for _ in range(n+1):
    #현재 가장 최단 거리가 짧은 노드를 꺼내서 방문 처리
    cur = get_smallest_node()
    visited[cur] = 1
    #현재 방문한 노드와 연결된 노드들에 대해 최단 거리 테이블 갱신
    for next in graph[cur]:
      #현재 노드까지의 최단 거리 + 현재 노드에서 다음 노드 까지의 거리
      cost = distance[cur] + next[1]
      #기존의 다음 노드까지의 최단 거리보다 현재 노드를 거친 다음 노드까지의 거리가 더 짧으면 테이블 갱신 
      if distance[next[0]] > cost:
        distance[next[0]] = cost

dijkstra(start)

for i in range(1, n+1):
  if i == INF:
    print("INFINITY")
  else:
    print(distance[i])
  

[이것이 취업을 위한 코딩 테스트다 with 파이썬 - 나동빈 (한빛미디어)] 교재를 참조하여 작성했습니다.