dijkstra 알고리즘 + lazy
다익스트라, 우선순위 큐와 게으른 삭제
Intro
다익스트라는 가중치가 있는 그래프에서 최단 경로를 찾는 알고리즘이다. 간선 비용이 전부 같다면 BFS로 충분하지만, 비용이 제각각이면 "먼저 도착한 게 더 가깝다"는 BFS의 전제가 깨진다.
그래서 다익스트라는 이 한 문장에 기댄다.
아직 확정되지 않은 정점들 중 거리 추정값이 가장 작은 정점은, 그 값이 곧 최단 거리다.
간선 비용이 음수만 아니라면, 다른 미확정 정점을 거쳐 돌아가는 길이 이보다 짧아질 수는 없다. 덕분에 한 번 확정한 정점은 다시 볼 필요가 없다. 그럼 남는 문제는 하나다. 최솟값을 매번 어떻게 빨리 찾느냐.
우선순위 큐
매번 전체를 훑어 최솟값을 찾으면 O(V²)가 된다. 정점이 십만 개쯤 되면 곤란하다.
필요한 기능은 "값 넣기"와 "가장 작은 값 꺼내기" 딱 두 개뿐이고, 이걸 O(log N)에 해주는 자료구조가 힙이다.
from heapq import heappush, heappop
def dijkstra(start, graph, N):
dist = [float('inf')] * N
dist[start] = 0
pq = [(0, start)]
while pq:
d, u = heappop(pq)
if d > dist[u]:
continue
for v, w in graph[u]:
nd = d + w
if nd < dist[v]:
dist[v] = nd
heappush(pq, (nd, v))
return dist
전체 시간은 O(E log V)로 줄어든다. 튜플은 반드시 (거리, 정점) 순서로 넣어야 한다. 뒤집으면 정점 번호로 정렬되어, 코드는 멀쩡히 도는데 답만 틀리는 버그가 난다.
게으른 삭제
눈에 잘 안 띄는 이 두 줄이 사실 핵심이다.
if d > dist[u]:
continue
원래 다익스트라는 큐에 든 항목의 값을 낮추는 decrease-key 연산을 쓴다. 그런데 heapq에는 그런 기능이 없다. 그래서 발상을 뒤집는다. 고칠 수 없다면 새로 하나 더 넣는다.
더 짧은 경로를 찾으면 기존 항목은 내버려 두고 새 항목만 push한다. 같은 정점이 큐에 여러 번 들어가지만, 힙은 작은 것부터 꺼내주므로 처음 나온 것만 유효하고 나머지는 철 지난 정보다. 그걸 걸러내는 게 위 두 줄이다. "꺼낸 값이 지금 아는 값보다 크네? 옛날 거네" 하고 넘기는 것이다.
이렇게 미리 지우지 않고 마주쳤을 때 버리는 방식을 게으른 삭제라고 부른다.
재미있는 건 이 두 줄을 지워도 답은 맞다는 점이다. 아래쪽 nd < dist[v]가 한 번 더 막아주기 때문이다. 대신 낡은 항목마다 이웃을 전부 다시 순회하느라 시간이 크게 늘어난다. 정답이 아니라 속도를 위한 줄이라, 빼먹으면 오답이 아니라 시간 초과로 돌아온다.
마무리
다익스트라 코드는 열 줄 남짓이라 외워서 쓰기 쉽다. 그래도 힙을 쓰는 이유는 "최솟값을 빨리 꺼내려고", 저 두 줄이 있는 이유는 "decrease-key가 없어서"라는 것 정도는 알아두면 좋다. 코드가 안 돌 때 어디를 봐야 할지 감이 잡힌다.
별거 아닌 것처럼 생겼지만, 사실 제일 열심히 일하고 있는 두 줄이다.