공부/알고리즘
다익스트라(Dijkstra) - Python, heapq
아찌방
2025. 2. 7. 23:40
코드
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