알고리즘

[python] dijkstra 알고리즘

ujin2021 2021. 3. 22. 21:50

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 들이 변하는지 보기위해서 넣은것이다.