다익스트라 알고리즘 수행 시 매 번 최단 거리가 가장 짧은 노드에 방문하여 해당 노드를 거쳐 가는 거리를 계산하여 최단 거리 테이블을 갱신한다.
이전 포스트에서는 get_smallest_node() 함수를 통해
매번 최단 거리 테이블을 선형 탐색하여 최소값을 구하는 방식으로 구현했다.
이 경우에 V 개의 노드에 방문할 때마다 V 개의 최단 거리를 탐색하므로 O(V2) 시간 복잡도인 것을 알 수 있었다.
이번에는 Heap 자료구조를 사용하여 시간복잡도가 개선된 다익스트라 알고리즘을 구현한다.
Heap은 Priorty queue (우선순위 큐)를 구현하기 위해 사용된다.
우선순위 큐는 FIFO 구조의 일반 큐와 다르게 삽입되는 데이터의 가치에 따라 먼저 삭제된다.
우선순위 큐를 최소 힙으로 구현하는 경우 가치가 낮은 데이터가 먼저 삭제되고
최대 힙으로 구현하는 경우 가치가 높은 데이터가 먼저 삭제된다.
예를 들어, 최소 힙으로 구현된 우선순위 큐의 아래와 같이 데이터가 삽입되는 경우,
- push((10, 'a'))
- push((4, 'b'))
- push((7, 'c'))
- push((2, 'd'))
삽입된 순서에 상관 없이 (2, 'd'), (4, 'b'), (7, 'c'), (10, 'a') 순서로 pop된다.
물론, 최대 힙으로 구현된 우선순위 큐의 경우 위와 반대 순서로 동작할 것이다.
개선된 다익스트라 알고리즘의 경우, 우선순위 큐에 (최단 거리, 노드) 튜플을 삽입하고 매번 가장 위에 있는 노드에 방문한다. 즉, 최단 거리가 가장 짧은 노드에 방문한다.
현재 방문한 노드를 거쳐 다음 노드로의 거리 계산 후 최단 거리 테이블이 갱신될 때,
해당 노드의 (최단 거리, 노드) 튜플을 삽입한다.
기존 getsmallest_node 함수에서 매번 distance 리스트를 선형 탐색했다면
heap에 데이터를 삽입하여 우선순위가 결정되는 데 걸리는 시간 복잡도는 O( log N)이기 때문에 최종 시간 복잡도는 O(V2)에서 O(E log V)로 대폭 감소된다.
여기서 V는 노드의 수, E는 간선의 수를 의미한다.
파이썬에서 우선순위 큐를 사용하기 위해서는 heapq 라이브러리를 사용하며 기본적으로 최소힙을 사용한다. 만약 최대 값을 먼저 삭제해야 하는 경우, -1을 곱한 음수 가치를 입력하고 데이터를 꺼낸 후 다시 -1을 곱해 양수로 만드는 스킬을 사용하기도 한다.
아래는 파이썬으로 구현된 코드이다.
import sys
import heapq
INF = int(1e9)
input = sys.stdin.readline
n, m = map(int, input().split())
start = int(input())
graph = [[] for _ in range(n+1)]
distance = [INF] * (n+1)
for _ in range(m):
a, b, c = map(int, input().split())
graph[a].append((b, c))
def dijkstra(start):
q = []
distance[start] = 0
heapq.heappush(q, (0, start))
while q:
dist, cur = heapq.heappop(q)
if distance[cur] < dist:
continue
for next in graph[cur]:
cost = dist + next[1]
if distance[next[0]] > cost:
distance[next[0]] = cost
heapq.heappush(q, (cost, next[0]))
dijkstra(1)
for i in range(1, n+1):
if i == INF:
print("INFINITY")
else:
print(distance[i])[이것이 취업을 위한 코딩 테스트다 with 파이썬 - 나동빈 (한빛미디어)] 교재를 참조하여 작성했습니다.
