본문 바로가기
공부/알고리즘

다익스트라(Dijkstra) - Python, heapq

by 아찌방 2025. 2. 7.

 

코드

import heapq
import sys

def dijkstra(graph, start):
    INF = float('inf')
    n = len(graph)
    distance = [INF] * n
    distance[start] = 0
    
    pq = []
    heapq.heappush(pq, (0, start))  # (거리, 정점)

    while pq:
        dist, now = heapq.heappop(pq)

        if distance[now] < dist:
            continue  # 이미 처리된 노드라면 무시

        for neighbor, weight in graph[now]:
            cost = dist + weight
            if cost < distance[neighbor]:  # 최단 거리 갱신
                distance[neighbor] = cost
                heapq.heappush(pq, (cost, neighbor))

    return distance

# 예제 그래프 (인접 리스트)
graph = {
    0: [(1, 4), (2, 1)],
    1: [(3, 1)],
    2: [(1, 2), (3, 5)],
    3: []
}

start_node = 0
shortest_distances = dijkstra(graph, start_node)
print(shortest_distances)  # 시작 노드(0)에서 각 정점까지의 최단 거리

 

과정

 

 

  • 거리 배열 초기화: 시작 정점에서 모든 정점까지의 최단 거리를 무한(∞)으로 설정, 시작 정점의 거리는 0으로 설정
  • 우선순위 큐(Priority Queue) 사용: 시작 정점을 큐에 넣음 (우선순위 큐는 Python의 heapq 사용)
  • 최단 거리 갱신:
    • 우선순위 큐에서 가장 최단 거리가 짧은 정점을 꺼냄
    • 해당 정점과 연결된 다른 정점들의 거리 갱신
    • 갱신된 거리가 기존보다 짧다면 큐에 추가
  • 모든 정점에 대한 거리 갱신이 끝나면 종료

 

시간 복잡도

  • 우선순위 큐(힙) 사용 시: O(E log V)
    • V: 정점 개수
    • E: 간선 개수
    • 우선순위 큐에서의 삽입/삭제 연산이 O(log V) 이므로, 모든 간선을 탐색하는 과정이 O(E log V) 이다.
  • 일반적인 구현 (배열 이용): O(V²)
    • 매번 최소 거리 정점을 찾는 데 O(V) 가 걸리므로, 전체적으로 O(V²) 이다. -> heap 사용해야하는 이유

 

728x90