[보기 쉬운 Algorithm] 다익스트라(Dijkstra) 알고리즘 - Python

'다익스트라(Dijkstra) 알고리즘'은 비용(가중치) 간선이 있는 노드 그래프의 최단 거리를 구하는 알고리즘

 

시간을 줄이기 위해서 파이썬 모듈 중 최소 힙을 구현할 수 있는 heapq를 사용한다.

↳ list로도 구현 가능하나 시간이 느려진다.

 

import heapq


def dijkstra(graph, start):
    distances = {node: float('inf')
                 for node in graph}
    distances[start] = 0
    queue = []
    heapq.heappush(queue, [distances[start], start])

 

graph와 시작 점인 start를 인자로 받는다.

distances는 계산된 각 점마다의 거리를 담고있다. 다익스트라 알고리즘은 거리가 짦은 값을 갱신하므로, graph에 있는 node의 개수 만큼 모든 노드의 거리를 가장 큰 수인 inf로 초기화 해준다.

 

 

queue에는 heapq를 사용하여, 우선순위가 되는 거리, 다음에 노드를 추가한다. [거리, 기준 노드]

 

    while queue:
            current_distance, current_node = heapq.heappop(queue)

            if distances[current_node] < current_distance:
                continue

            for new_node, new_distance in graph[current_node].items():
                distance = current_distance + new_distance
                if distance < distances[new_node]:
                    distances[new_node] = distance

                    heapq.heappush(queue, [distance, new_node])

        return distances

 

queue가 사라질 때 까지 while문을 반복한다.

queue에 있는 현재 노드의 거리와, 노드를 받아온다.

만약 distances에 담겨있는 거리보다 길다면 해당 거리와 노드는 무시한다. (최단 거리를 찾기 때문에)

 

graph에서 목적지 node와 그 node로 가는 비용을 받아와서 현재 노드와 더해 총 목적지까지의 총 거리를 distance에 구한다. 만약 distance가 distacnes[new_node]에 있는 값보다 짧을 경우 값을 갱신하고, 그 노드가 기준이 될 수 있게 queue에 집어 넣는다.

 

while문 반복이 끝나면 distances를 반환한다.

 

<전체 코드>

import heapq


def dijkstra(graph, start):
    distances = {node: float('inf')
                 for node in graph}
    distances[start] = 0
    queue = []
    heapq.heappush(queue, [distances[start], start])

    while queue:
        current_distance, current_node = heapq.heappop(queue)

        if distances[current_node] < current_distance:
            continue

        for new_node, new_distance in graph[current_node].items():
            distance = current_distance + new_distance
            if distance < distances[new_node]:
                distances[new_node] = distance

                heapq.heappush(queue, [distance, new_node])

    return distances

<예시 입력>

graph = {
    'A': {'B': 2, 'C': 4},
    'B': {'C': 3, 'D': 8},
    'C': {'D': 2},
    'D': {},
    'E': {'A': 8},
}
print(dijkstra(graph, 'A'))

<예시 출력>

{'A': 0, 'B': 2, 'C': 4, 'D': 6, 'E': inf}