📌다익스트라 최단 경로 알고리즘 (이코테 WITH 파이썬)

이번달 코딩테스트 중에서 이거 딱 다익스트라인데 하면서 못 푼 문제가 있었다. 끝나고 후기를 찾아보니 정말 다익스트라로 풀 수 있는 문제여서 아쉬웠던 기억이 있다.

 

이참에 다익스트라 알고리즘을 복습해보자.

 

최단 경로 알고리즘

말 그대로 가장 짧은 경로를 찾는 알고리즘. 그래프를 이용해 표현하며 각 지점은 노드로, 지점간의 연결은 간선으로 표현한다. 

 

다익스트라 알고리즘

다익스트라 알고리즘은 그래프의 여러개의 노드가 있을 때 특정한 노드에서 출발해 다른 노드로 가는 각각의 최단경로를 구하는 알고리즘이다.

 

과정은 다음과 같다.

 

출발 노드 설정 > 최단 거리 테이블 초기화 > 방문하지 않은 노드 중 최단 거리가 가장 짧은 노드를 선택 > 해당 노드에서 부터 다른 노드로 가는 거리를 계산한 후 테이블 갱신 > 반복

 

방문하지 않는 노드 중 최단 거리가 가장 짧은 노드를 선택한다는 점에서 그리디 알고리즘이라고도 할 수 있다.

 

 

💡이때 최단 거리가 가장 짧은 노드를 매번 탐색하고, 탐색한 노드와 연결된 노드를 일일이 확인하면 시간복잡도가 O(V^2)이 된다. 시간복잡도를 줄이기 위해 최단 거리가 가장 짧은 노드를 선형적으로 탐색하는 것이 아닌, 힙 자료구조를 사용할 수 있다.

* 힙 자료구조는 완전 이진 트리의 일종으로 삽입 삭제 시간복잡도가 O(logN)이다.

 

힙 자료구조는 우선순위 큐(우선순위가 가장 높은 데이터 부터 추출)를 구현하기 위해 사용하는 자료구조이다. 파이썬에서는 heapq를 사용할 수 있다. 파이썬 라이브러리에서는 기본적으로 최소 힙 구조를 사용하기 때문에 최단 거리가 가장 짧은 노드를 먼저 탐색하는 다익스트라 알고리즘에 그대로 사용하면 된다.

 

코드 구현

그래프를 표현할 때에는 리스트를 노드의 개수 + 1로 할당하는 습관을 들이자(노드 번호로 리스트에 바로 접근).  

import heapq
import sys
input = sys.stdin.readline
INF = int(1e9)  # 무한을 의미하는 값 10억 설정

# 노드와 간선의 개수
n, m = map(int, input().split())
# 시작 노드    
start = int(input())

# 각 노드의 연결 정보를 담을 리스트
graph = [[] for i in range(n + 1)]
distance  = [INF] * (n + 1 )

# 모든 간선 정보 입력 
for _ in range(m):
    a, b, c = map(int, input().split())
    # a번 노드에서 b번 노드로 가는 비용 c
    graph[a].append((b, c)) 

def dijkstra(start):
    q = []
    # 시작 노드로 가는 최단 경로는 0
    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, n + 1):
    # 도달할 수 없는 경우
    if distance[i] == INF:
        print('INFINITY')
    else:
        print(distance[i])

 

 

---

참고자료

(이코테 2021) 이것이 취업을 위한 코딩 테스트다 with 파이썬