다익스트라(dijkstra) 알고리즘 (feat. heapq)

하나의 정점으로 부터 모든 다른 정점까지의 최단 경로를 찾는 최단 경로 알고리즘인 다익스트라 알고리즘에 대해 공부해보고자 한다.

 

최단 경로 알고리즘의 아이디어
1. 출발 노드와 도착 노드 설정
2. 알고 있는 모든 거리 값 부여
3. 출발 노드부터 시작해 방문하지 않은 인접 노드 방문, 거리 계산하고 현재 알고있는 거리보다 짧으면 값 갱신
4. 현재 노드에 인접한 모든 미방문 노드까지 거리 계산했으면, 현재 노드를 미방문 집합에서 제거 5. 도착 노드가 미방문 노드 집합에서 벗어나면, 알고리즘 종료

 

사전 배경 지식

heapq 모듈 자료구조
heapq 모듈은 이진 트리 기반의 최소 힙 자료구조를 제공한다.
최소 힙을 사용하면 원소들이 항상 정렬된 상태로 추가되고 삭제되며, 최소 힙에서 가장 작은 값은 언제나 인덱스 0, 즉, 이진 트리의 루트에 위치한다.
내부적으로 최소 힙 내의 모든 원소는 항상 자식 원소들 (2k+1, 2k+2)보다 크기가 작거나 같도록 원소가 추가되고 삭제된다.

(min heap 자료구조 예시)

     1  ---> root
   /   \
  3     5
 / \   /
4   8 7

 

# 모듈 임포트
from heapq import heappush, heappop

# 최소 힙 생성
# 파이썬에서는 heapq 모듈을 통해서 원소를 추가하거나 삭제한 리스트가 그냥 최소 힙이다.
heap = []

# 힙에 원소 추가
from heapq import heappush
heappush(heap, 4)
heappush(heap, 1)
heappush(heap, 7)
heappush(heap, 3)
print(heap)

>>> [1, 3, 7, 4]

# 힙에서 원소 삭제
from heapq import heqppop
print(heappop(heap))
print(heap)
>>> 1
>>> [3, 4, 7]

# 기본 리스트를 힙으로 변환
from heapq import heapify
heap = [4, 1, 7, 3, 8, 5]
heapify(heap)
print(heap)
>>> [1, 3, 5, 4, 8, 7]

코드 구현

import heapq

def dijkstra(graph, start):
    distances = {node: float('inf') for node in graph}
    distances[start] = 0 # 시작 값은 0
    queue = []
    heapq.heappush(queue, [distances[start], start]) # 시작 노드부터 탐색 시작
    
    while queue:
        current_distance, current_destination = heapq.heappop(queue) # 탐색 할 노드, 거리를 가져온다.
        
        if distances[current_destination] < current_distance: # 기존에 있는 거리보다 길다면, 넘어감
            continue
        
        for new_destination, new_distance in graph[current_destination].items():
            distance = current_distance + new_distance # 해당 노드를 거쳐 갈 때 거리
            if distance < distances[new_destination]: # 알고 있는 거리보다 작으면 갱신
                distances[new_destination] = distance
                heapq.heappush(queue, [distance, new_destination]) # 다음 인접 거리를 계산 하기 위해 큐에 삽입
    return distances

in

print(dijkstra(graph, 'A'))

out

{'A': 0, 'B': 6, 'C': 1, 'D': 2, 'E': 5, 'F': 6}

 

참고 : https://justkode.kr/algorithm/python-dijkstra