[알고리즘] 다익스트라(Dijkstra Algorithm) 알고리즘 파이썬, 최소힙(min-heap) 사용하기

# 파이썬 언어

# 무한대값 만들기
CONST_INFINITE = float('inf')

import heapq

# heapq라이브러리를 사용하여 최소 힙 만들기
min_heap = []

# push O(log n)
heapq.heappush(min_heap,[10,'A'])

# pop O(1)
heapq.heappop(min_heap)

 heapq에서는 최소힙만 지원하기 때문에, max_heap을 사용하기 위해서는 값을 insert할때 ' - '을 곱해주고 꺼낼때 다시 ' - '를 곱하는 편법을 이용한다고 한다. 뭐 우선순위 큐에서 람다식으로 우선순위를 정해주는 방법도 있을거 같기는 하다. 

# 사용된 그래프

# 설명

 학교에서 자료구조시간에 다익스트라 알고리즘을 배울때는, 힙을 사용하지 않는 방법에 대해서 공부했었다. 배열을 읽고 배열에서 가장 작은 노드를 뽑아서 구현하는 방식을 사용했었는데, 힙을 사용하니 편하게 구현 할 수 있었다.

 BFS와 유사한 방식으로 진행되어 나간다. 힙에 해당 노드까지의 최신화 된 거리(weight)를 기록해 둔 후, 하나씩 pop하고 인접노드들 + 현재 거리의 값을 현재 distance_state와 비교해서 작은 update를 시켜주는 방식을 사용한다. 

 

# 코드

import heapq

CONST_INFINITE = float('inf')

def make_graph():
  graph = {
    'A': {'B': 8, 'C': 1, 'D': 2},
    'B': {},
    'C': {'B': 5, 'D' : 2},
    'D': {'E': 3, 'F': 5},
    'E': {'F': 1},
    'F': {'A': 5}
    }

  return graph

def dijstra(graph, start):
  min_queue = []
  distance_state = {node: CONST_INFINITE for node in graph}

  # 초기화 해주기
  distance_state[start] = 0
  heapq.heappush(min_queue, [distance_state[start], start])
  print(f"__init__ {distance_state}")

  # 힙에서 꺼내오고 비교하기
  while min_queue:
    cur_distance, cur_node = heapq.heappop(min_queue)  # ex) [ 0 : 'A']
    # 이미 최단경로이면 검사할 필요가 없음
    if cur_distance > distance_state[cur_node] :
      continue

    for adjacent_node, weight in graph[cur_node].items() :
      if distance_state[adjacent_node] > cur_distance + weight :
        # update
        update_weight = cur_distance + weight
        distance_state[adjacent_node] = update_weight
        heapq.heappush(min_queue,[update_weight,adjacent_node])
    print(f"distance_state = {distance_state} cur_node : {cur_node}")

  return distance_state


if __name__ == '__main__':
  start = 'A'
  data = make_graph()
  result = dijstra(data,start)

  print(f"\nresult : {result}")

 

# 참고

E = Edge

일단 모든 간선을 하나씩 순회하면서 탐색해야하기 때문에 O(E)의 시간이 걸린다.

힙을 insert할때 O(log n)만큼의 시간이 걸리는것을 생각해 봤을때를 생각해야 하기때문에 시간복잡도는 아래와 같다