다익스트라 알고리즘
Dijkstra 알고리즘은 그래프에서 최단 거리를 구하는 문제를 푸는 대표적인 알고리즘이다. 기본 알고리즘은 $O(N^2)$ 시간 복잡도를 가지지만, 우선순위 큐를 이용해 $O(NlogN)$로 최적화할 수 있다.
문제
N개의 도시가 있다. 그리고 한 도시에서 출발하여 다른 도시에 도착하는 M개의 버스가 있다. 우리는 A번째 도시에서 B번째 도시까지 가는데 드는 버스 비용을 최소화시키려고 한다. A번째 도시에서 B번째 도시까지 가는데 드는 최소비용을 출력하여라. 도시의 번호는 1부터 N까지이다.
- 백준 1916
문제를 도식화하면 위와 같다. 1번 도시에서 출발해 5번 도시에 도착한다. 이때 가장 짧은 거리로 이동하려 한다.
풀이
다익스트라 알고리즘은 생각보다 간단하게 작동한다.
- 출발점으로부터 거리를 모두 inf(무한)로 둔다. 자기 자신은 0으로 둔다.
- 인접한 도시를 모두 방문하고 최단 거리를 업데이트한다.
- 현재 상태에서 최단 거리가 가장 짧은 도시를 방문한다.
- 더 이상 방문할 도시가 없을 때까지 2 ~ 3을 반복한다.
한 단계씩 그림을 살펴보며 알아보자.
1. 출발점으로부터 거리를 모두 inf(무한)로 둔다. 자기 자신은 0으로 둔다.
1번 도시에서 1번 도시로 가는 거리는 0이다. 그리고 다른 도시로 가는 최단 거리는 아직 모르기 때문에 아주 큰 값인 inf로 둔다.
2. 인접한 도시를 방문하고 최단 거리를 업데이트한다.
비고: 이미 알고 있는 최단 거리(inf)와 새로 발견한 길을 비교해 짧은 쪽 데이터를 저장한다.
3. 현재 상태에서 최단 거리가 가장 짧은 도시를 방문한다.
최단 거리로 저장된 0km, 2km, 3km, 1km, 10km 중 0km가 가장 짧다. 하지만 1번 도시는 이미 방문했기 때문에 다시 방문하지 않아도 된다. 따라서 다음으로 짧은 1km인 4번 도시를 방문한다.
4. 현재 도시에 인접한 도시를 방문하고 최단 거리를 업데이트한다.
1번 → 4번으로 가는데 이미 1km가 필요하다. 따라서 1번 → 4번 → 5번은 { 1 + 3 }으로 4km가 된다. 기존에 알고 있던 10km(1번 → 5번)보다 짧으므로 최단 거리를 업데이트해준다.
5. 현재 상태에서 최단 거리가 가장 짧은 도시를 방문한다.
6. 현재 도시에 인접한 도시를 방문하고 최단 거리를 업데이트한다.
7. 현재 상태에서 최단 거리가 가장 짧은 도시를 방문한다.
8. 현재 도시에 인접한 도시를 방문하고 최단 거리를 업데이트한다.
마찬가지로 출발점 → 3번까지 이미 3km가 걸리기 때문에 { 3 + 거리 }로 새로운 거리를 계산한다.
9. 현재 상태에서 최단 거리가 가장 짧은 도시를 방문한다.
10. 더 이상 방문할 도시가 없을 때 종료한다.
11. 결과
1번 → 5번까지 최소 4km를 이동해야 한다는 결론을 얻었다. 과정을 역추적해보면 1번 → 4번 → 5번을 통해 이동했음을 알 수 있다.
용어 정리
코드를 이해하기 위해 기본적인 용어를 정리하고 시작하겠다.
그래프
vertex: 정점, edge: 간선, weight: 가중치
vertex는 그래프의 정점으로 문제에서는 각 도시에 해당한다. edge는 간선으로 vertex를 연결하는 선이다. weight는 각 edge가 가지는 가중치다. 문제에서는 거리를 의미한다.
heapq
heapq는 우선순위 큐로 pop 했을 때 가장 작은 값을 꺼내준다.
import heapq
data = []
heapq.heappush(data, 50)
heapq.heappush(data, 10)
heapq.heappush(data, 30)
print(heapq.heappop(data)) # 10
print(heapq.heappop(data)) # 30
print(heapq.heappop(data)) # 50리스트에 50, 10, 30 순서로 push 했지만 10, 30, 50 순서로 pop 된다. heapq를 이용하면 최단 거리를 가지는 도시를 효율적으로 찾을 수 있다.
Python 구현
일반 구현
import sys
def dijkstra(start):
# 시작 정점에 대해서 초기화
distance[start] = 0
visited[start] = True
for v, weight in graph[start]:
if distance[v] > weight:
distance[v] = weight
# 시작 정점를 제외한 전체 정점
for _ in range(num_vertex - 1):
current = get_closest()
visited[current] = True
update_distance(current)
def get_closest():
min_distance = INF + 1
closest_vertex = None
for vertex in range(1, num_vertex + 1):
if distance[vertex] < min_distance and not visited[vertex]:
min_distance = distance[vertex]
closest_vertex = vertex
return closest_vertex
def update_distance(vertex):
for adj_vertext, adj_weight in graph[vertex]:
cost = distance[vertex] + adj_weight
if cost < distance[adj_vertext]:
distance[adj_vertext] = cost
## 입력
input = sys.stdin.readline
num_vertex = int(input())
num_edge = int(input())
INF = 1e8
# vertex가 1부터 시작하기 때문에 range+1
graph = [[] for _ in range(num_vertex + 1)]
visited = [False for _ in range(num_vertex + 1)]
distance = [INF for _ in range(num_vertex + 1)]
for _ in range(num_edge):
u, v, weight = map(int, input().split(" "))
graph[u].append((v, weight))
start, target = map(int, input().split(" "))
## 출력
dijkstra(start)
print(distance[target])dijkstra 함수를 보면 위에서 소개한 방식과 동일하게 구현되었다.
- 출발점으로부터 거리를 초기화한다.
- 인접한 도시를 모두 방문하고 최단 거리를 업데이트한다.
- 현재 상태에서 최단 거리가 가장 짧은 도시를 방문한다.
- 더 이상 방문할 도시가 없을 때까지 반복한다.
그럼 주의해야 하는 반례를 살펴보자.
주의해야 하는 반례
1. inf 범위
문제 조건을 살펴보면 "버스 비용은 0보다 크거나 같고, 100,000보다 작은 정수이다."라는 조건이 있다. 따라서 inf를 100,000으로 설정하면 오류가 발생한다.
시작점부터 축적되어 온 가중치는 100,000보다 클 수 있다.
2. 길이 여러 개일 때
이 조건을 무시하고 실행해 봤더니 통과하지 못했다. 문제 중 여러 간선을 가진 케이스가 있나 보다.
우선순위 큐 활용
이번에는 heapq를 이용해 구현해 봤다.
import sys
import heapq
def dijkstra(start, target):
weights = [float("inf") for _ in range(num_vertex + 1)]
weights[start] = 0
will_visit = [(weights[start], start)] # 방문할 도시 (우선순위 큐)
while will_visit:
# 현재 상태에서 최단 거리가 가장 짧은 도시를 방문
# 방문한 노드는 pop으로 제거
weight, vertex = heapq.heappop(will_visit)
if weights[vertex] < weight:
# 최단거리보다 멀 때는 무시
continue
for adj_vertex, adj_weight in graph[vertex]:
cost = weight + adj_weight # 새로운 거리
if cost < weights[adj_vertex]:
# 기존 거리보다 짧으면 업데이트
weights[adj_vertex] = cost
heapq.heappush(will_visit, (cost, adj_vertex))
return weights[target]
## 입력
input = sys.stdin.readline
num_vertex = int(input())
num_edge = int(input())
graph = [[] for _ in range(num_vertex + 1)]
for _ in range(num_edge):
u, v, weight = map(int, input().split())
graph[u].append((v, weight))
start, target = map(int, input().split())
## 출력
print(dijkstra(start, target))전체적으로 코드가 깔끔해졌다. 현재 상태에서 가중치가 가장 작은 정점을 찾는 일을 heapq가 알아서 해주기 때문이다.
- 일반 구현: 188ms, 116048KB
- heapq 구현: 272ms, 121228KB
속도 면에서는 오히려 heapq 구현이 더 느린 모습이다. ...왜?