[Python] 우선순위 큐를 활용한 개선된 다익스트라(Dijkstra) 알고리즘

🔊 이번 포스팅에는 최근에 Python으로 알고리즘을 공부하기 시작하면서 알게 된 여러 알고리즘의 원리와 Python으로 구현하는 방법에 대해 소개해보려 한다. 필자는 최근 알고리즘 공부를 '나동빈'님의 '이것이 취업을 위한 코딩 테스트다 with 파이썬'이라는 책과 백준 온라인 저지 사이트로 하고 있다. 이 중 '나동빈'님이 저자이신 책에서 가르쳐주는 내용을 기반으로 배운 내용을 정리해보려 한다.

 

이번 포스팅에서는 저번 시간에 배웠던 다익스트라 알고리즘의 간단한 구현 방법에 이어 이를 개선한 우선순위 큐 즉, heapq를 활용한 개선된 다익스트라 알고리즘에 대해 알아보려고 한다. 다익스트라의 구현 개념에 대해서는 거의 유사하므로 해당 포스팅에서는 우선순위 큐를 어떤 부분에 활용하고 Python으로 구현할 때 어떤 부분이 간소화되어 시간 복잡도가 기존에 비해 확 줄어드는지를 중심으로 알아보자.

 

가장 우선적으로 보내야 하는 차량을 어떻게 처리할까?

---

1. 힙? 우선순위 큐?

힙과 우선순위 큐의 관계는 다음과 같이 정의할 수 있다. 힙이라는 자료구조를 사용해 우선순위 큐를 구현한다. 그리고 구현된 우선순위 큐는 Python에서 PriorityQueue, heapq 라이브러리 형태로 제공되며, 이 우선순위 큐를 가지고 다익스트라 알고리즘의 병목현상을 개선한다. 이 때, 병목현상이란 저번 포스팅에서 언급했지만 기존의 다익스트라 알고리즘의 문제점 중 하나는 알고리즘 반복 단계에서 시작 노드와 가장 거리가 짧은 최단 거리 노드를 찾아야 하는데, 이 때 매번 선형탐색을 수행했어야 했다는 것이다. 

 

그렇다면 우선순위 큐는 무엇일까? 힙이라는 자료구조로 구현된다. 이렇게 힙으로 구현된 우선순위 큐는 어떤 장점이 있을까? 우리가 기존에 배웠던, DFS에서 사용되던 스택은 FILO(First-In-Last-Out) 방식을, BFS에서 사용되던 큐는 FIFO(First-In-First-Out) 방식으로 구현된다. 반면에 우선순위 큐는 가장 우선순위가 높은 데이터부터 Out(삭제)되는 방식을 취한다. 이 때 대부분의 프로그래밍 언어에서는 데이터의 우선순위를 결정하는 기준은 데이터의 가장 첫번째 원소값을 기준으로 한다. 결국 우선순위 큐는 데이터를 특정한 우선순위에 따라 처리하고 싶을 때 자주 사용된다.

 

우선순위 큐를 구현하는 여러가지 프로그래밍 언어의 라이브러리들은 최소 힙 또는 최대 힙을 구현한다. 최소 힙이란, 데이터의 값이 가장 낮은 것을 가장 우선으로 여겨 정렬하는 반면 최대 힙은 데이터의 값이 가장 큰 것을 가장 우선으로 여겨 정렬하게 된다. Python에서 제공하는 라이브러리들은 최소 힙 방식으로 제공된다. 보통 최단 경로 알고리즘은 간선 비용을 '최소'로 하는 것을 우선순위 기준으로 하기 때문에 Python을 활용하게 되면 heapq 또는 PriorityQueue를 변형하지 않고 그대로 사용해도 무방하다.

 

참고로 우선순위 큐를 구현하는 방법으로 힙과 리스트를 사용할 수 있는데, 리스트를 활용하게 되면 삭제 시 최악의 경우 시간 복잡도 $O(N)$이 걸리는 반면, 힙을 사용하게 되면 최악의 경우라도 삭제 시 시간 복잡도 $O(logN)$이 걸린다. 힙에 대한 자료구조의 설명은 여기를 참고해보자. 또 Python을 활용해 우선순위 큐(최소 힙 또는 최대 힙)를 구현하는 방법은 기존 포스팅을 참고하자.

2. 다익스트라 알고리즘에 우선순위 큐 활용하기

이제 우선순위 큐를 다익스트라 알고리즘에 어떻게 사용하는지 살펴보자. 다익스트라 알고리즘에서 우선순위 가준은 뭘까? 바로 시작 노드와의 거리가 가장 가까운 노드를 의미한다. 따라서 알고리즘을 반복할 때마다 실시간으로 시작노드와 가장 가까운 노드를 저장하기 위한 목적으로 우선순위 큐를 활용한다.

 

참고로 기존의 간단한 다익스트라 알고리즘을 구현할 때는 사전에 2가지 종류의 리스트를 먼저 정의하고 진행했다. 첫 번째는 최단 거리를 기록하면서 갱신할 거리 테이블, 두 번째는 방문한 노드인지 처리하고 확인하는 목적의 방문 테이블이였다. 하지만 우선순위 큐를 이용하게 되면 최단 거리를 기록할 거리 테이블만 정의한다. 왜냐하면 우선순위 큐를 이용하게 되면 우선순위 큐가 알아서 가장 최단 거리인 노드를 가장 앞으로 정렬해주기 때문에 기존에 기록한 최단 거리보다 크다면 그냥 무시해주면 된다. 만약 기존에 기록한 최단 거리보다 더 짧은 새로운 최단 거리가 등장하게 되면 해당 거리와 노드를 우선순위 큐에 넣어준다. 자, 이제 그렇다면 우선순위 큐를 활용해서 다익스트라 알고리즘이 어떻게 동작하는지 도식화해서 살펴보자.

2-1. 첫번째 초기 단계

초기 단계

 

시작 노드를 1이라고 가정하자. 시작 노드 1에서 1로 가는 거리는 없으므로 0이다. 그리고 시작 노드인 1번에서 1번 노드로 가는 거리는 0 하나이므로 하나의 데이터를 우선순위 큐에 넣어준다.

2-2. 두 번째 단계

이제 우선순위 큐에 있는 데이터인 (거리:0, 노드:1)을 뽑아서 1번 노드를 탐색하자. 

 

두 번째 단계

 

위에서 보다 시피 1번 노드와 연결되어 있는 노드들은 2,3,4번 노드들이다. 해당 노드들을 우선순위 큐에 다 집어넣으면 위와 같이 '거리'를 우선순위로 하여 자동 정렬된다. 위 예시에서는 시작노드 1번과 거리가 1로 가장 가까운 4번 노드 데이터가 우선순위 큐안에 가장 앞에 위치한다.(이렇게 해서 시작 노드와 가장 가까운 거리에 있는 노드를 탐색하는 과정을 우선순위 큐가 대신해준다!)

2-3. 세 번째 단계

다음은 우선순위 큐에 가장 앞에 있는 (거리:1, 노드:4) 데이터를 뽑아서 4번 노드를 탐색한다.

 

세번째 단계

 

4번 노드와 연결된 노드들은 3,5번노드들이다. 이 때, 중요한 포인트는 연결된 3,5번 노드들과 시작 노드들간의 거리가 이전보다 최단 거리로 업데이트될 경우에만 우선순위 큐에 집어넣는다. 예를 들어 3번 노드는 위 그림에서 보다시피 두 번째 단계에서는 거리가 5였지만 4번 노드를 탐색함으로써 더 최단거리인 거리가 4로 업데이트되었다. 즉, 이럴 때만 우선순위 큐에 집어넣어야 한다는 것이다. 그리고 우선순위 큐 안에서는 거리를 기준으로 가장 가까운 데이터를 다시 앞으로 자동 정렬시켜준다. 

 

방금 설명한 경우와는 반대로, 우선순위 큐에서 꺼낸 데이터의 거리가 기존의 거리 테이블의 거리값보다 클 경우도 있을 수 있다. 이럴 때는 그냥 무시하고 넘어가면 된다. 따라서 우선순위 큐에서 빼낸 노드를 탐색하려 하는데, 해당 노드보다 더 짧은 거리가 거리 테이블에 이미 갱신되어 있을 경우, 이 경우에 대해서는 무시해주고 넘어가면 된다.

3. Python으로 개선된 다익스트라 알고리즘 구현하기

이제 그렇다면 Python을 활용해 우선순위 큐를 활용한 다익스트라 알고리즘을 구현해보는 소스코드를 살펴보자.

 

import sys
import heapq
input = sys.stdin.readline
n, m = map(int, input().split())
start = int(input())
INF = int(1e9)
distance = [INF] * (n+1)
graph = [[] for _ in range(n+1)]
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, node = heapq.heappop(q)
        # 큐에서 뽑아낸 거리가 이미 갱신된 거리보다 클 경우(=방문한 셈) 무시
        if distance[node] < dist:
            continue
        # 큐에서 뽑아낸 노드와 연결된 인접노드들 탐색
        for next in graph[node]:
            cost = distance[node] + next[1]   # 시작->node거리 + node->node의인접노드 거리
            if cost < distance[next[0]]:      # cost < 시작->node의인접노드 거리
                distance[next[0]] = cost
                heapq.heappush(q, (cost, next[0]))


dijkstra(start)

for i in range(1, len(distance)):
    if distance[i] == INF:
        print('도달할 수 없음')
    else:
        print(distance[i])

4. 개선된 다익스트라 알고리즘의 시간 복잡도

지금까지 배운 우선순위 큐를 활용한 다익스트라 알고리즘의 시간 복잡도는 $O(ElogV)$이다. 이 때, $E$는 간선의 개수, $V$는 노드의 개수를 의미한다. 왜 간선의 개수 $E$가 시간 복잡도에 관여할까? 위 소스코드를 보면 우선순위 큐에서 꺼낸 노드를 탐색할 때 해당 노드와 연결된 노드들만 탐색하므로 최악의 경우라도 총 간선의 개수인 $E$만큼을 반복하게 된다. 따라서 우선순위 큐를 활용한 다익스트라 알고리즘은 $E$개의 간선을 힙에 넣었다가 다시 빼는 것으로 볼 수 있으므로 $O(ElogE)$가 되게 된다. 이 때, $E$는 $V^2$보다 항상 작으므로 $logE$는 $log{V^2}$보다 작다. 따라서 하나의 간선에 대해 걸리는 시간 복잡도는 $O(logV)$라고 볼 수 있으며 최악의 경우 최대 간선 개수를 고려해 총 시간 복잡도는 $O(ElogV)$가 된다.