본문 바로가기
알고리즘

[python] dijkstra 알고리즘

by ujin2021 2021. 3. 22.

dijkstra(다익스트라) 알고리즘은 source node로 부터 모든 node까지의 최단거리를 구하는 알고리즘이다. 

먼저 다른 블로그에서 코드를 보고, 출력을 해보며 작동원리를 이해해보았다.

 

다익스트라 알고리즘을 구현하기전에 먼저 heapq에 대해 공부해야 했다. heapq는 우선순위큐 를 구현하기 위해 쓰이는 모듈이다. 따로 설치할 필요없이 import heapq 하면 된다.

 

[heapq]

  • 완전 이진 트리 형태
  • 부모 노드 키값이 자식 노드 키값보다 항상 작다(최소힙)
  • 키값의 대소관계는 부모/자식 관계에서만 성립하고, 자식끼리의 대소관계는 없다

heapq 예시

import heapq

q = []
heapq.heappush(q, [2, 'D'])
heapq.heappush(q, [8, 'B'])
heapq.heappush(q, [3, 'E'])
heapq.heappush(q, [1, 'C'])
heapq.heappush(q, [6, 'F'])
print(q) # [[1, 'C'], [2, 'D'], [3, 'E'], [8, 'B'], [6, 'F']]

 

사실 처음에 heapq 원리를 이해하지 못하고 다익스트라 알고리즘을 먼저 보았다가 이해가 안되는 부분이 생겼다. 바로 정렬되는 순서때문인데, 위의 heapq 특징 세번째를 기억하면 된다(키값의 대소관계는 부모/자식 관계에서만 성립하고, 자식끼리의 대소관계는 없다)

 

위의 코드 설명을 하자면, 먼저 heapq 모듈을 import해주고, 빈 리스트 q를 만들어 준다. heapq.heappush(리스트, 넣을 값) 으로 q에 값을 넣을 수 있다.

 

 

[다익스트라 알고리즘]

import heapq # heapq : 우선순위 큐 구현

def dijkstra(graph, start) :
    distances = {node: float('inf') for node in graph} # start로 부터 노드까지의 거리(inf로 초기화)
    distances[start] = 0 # start -> start 니까 0으로 설정
    queue = []
    heapq.heappush(queue, [distances[start], start]) # queue = [0, 'A'] (시작노드부터 탐색)
    
    while queue : # queue에 남아있는 노드가 없으면 끝난다
        print('\n new while')
        current_distance, current_destination = heapq.heappop(queue) # 거리, 탐색할 노드를 가지고 온다
        print(f'A(start) -> {current_destination}(current) : {current_distance}')
        # current_distance : start -> destination 거리, current_destination : 탐색할 노드

        print(f'distances[current_destination] : {distances[current_destination]} < current_distance : {current_distance} => {distances[current_destination] < current_distance}')
        if distances[current_destination] < current_distance :
            # 계산해 놓은 A -> destination 거리보다 현재 거리가 더 크면 볼필요 없음
            print('continue')
            continue
        # 만약 계산해놓은 거리보다 지금이 같거나 짧다면 => 짧은걸로 갱신해야겠지
        print('items : ',graph[current_destination].items())
        for new_destination, new_distance in graph[current_destination].items():
            print(f'[new_destination, new_distance] : {new_destination}, {new_distance}')
            distance = current_distance + new_distance # 해당 노드를 거쳐 갈 때 거리
            print(f'distance : {distance} (current_distance : {current_distance} + new_distance : {new_distance})')
            print(f'distance({distance}) < distances[new_destination]({distances[new_destination]}) => {distance < distances[new_destination]}')
            if (distance < distances[new_destination]) : # 알고 있는 거리 보다 작으면 갱신
                print('===renew distance===')
                distances[new_destination] = distance
                heapq.heappush(queue, [distance, new_destination]) # 다음 인접 거리를 계산하기 위해 큐에 삽입
                print(f'after renew distance queue : {queue}')
    return distances
 
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}
}
print(dijkstra(graph, 'A'))

 

print문은 어떤식으로 queue와 distances 들이 변하는지 보기위해서 넣은것이다.