구현

이 장에서는 Python으로 이진 힙 다익스트라와 유니온 파인드 크루스칼을 만들고, 같은 다익스트라를 C++, Java, TypeScript로 옮깁니다. 정점 번호는 0..n-1이고 그래프는 인접 리스트 adj[u] = [(v, w), ...]입니다.

Python 다익스트라

import heapq
from math import inf

def dijkstra(adj, source):
    n = len(adj)
    dist = [inf] * n
    parent = [-1] * n
    dist[source] = 0
    heap = [(0, source)]
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist[u]:  # 낡은 항목: u는 이미 더 짧은 거리로 확정됨
            continue
        for v, w in adj[u]:
            nd = d + w
            if nd < dist[v]:
                dist[v] = nd
                parent[v] = u
                heapq.heappush(heap, (nd, v))
    return dist, parent

def path_to(parent, target):
    path = []
    while target != -1:
        path.append(target)
        target = parent[target]
    return path[::-1]

adj = [[(1, 4), (2, 1)], [(3, 1)], [(1, 2), (3, 5)], [(4, 3)], []]
dist, parent = dijkstra(adj, 0)
print(dist)  # [0, 3, 1, 4, 7]
print(path_to(parent, 4))  # [0, 2, 1, 3, 4]

한 줄씩 보면 다음과 같습니다.

• heapq는 일반 리스트를 이진 최소 힙으로 다룹니다. 튜플은 첫 항목부터 비교하므로 거리가 가장 작은 쌍이 먼저 나옵니다.
• dist는 무한대로, parent는 -1(아직 이전 정점 없음)로 시작합니다.
• if d > dist[u]: continue가 지연 삭제 검사입니다. 이 줄이 없으면 낡은 쌍마다 u의 간선을 다시 훑어 실행 시간이 늘어납니다.
• 거리가 줄면 dist[v]와 parent[v]를 고치고 새 쌍을 넣습니다. 옛 쌍은 힙에 남아 있다가 꺼낼 때 건너뜁니다.
• path_to는 parent를 따라 출발점까지 거슬러 올라갑니다. 닿을 수 없는 목적지라면 [target]만 돌려주므로 먼저 dist[target] != inf를 확인합니다.

출발점과 목적지가 하나씩인 질의라면 목적지를 꺼낸 순간 바로 돌려주면 됩니다. 그 시점에 거리가 확정되기 때문입니다.

Python 크루스칼과 유니온 파인드

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]  # 경로 절반 줄이기
            x = self.parent[x]
        return x

    def union(self, a, b):
        a, b = self.find(a), self.find(b)
        if a == b:
            return False  # 합치면 사이클이 생김
        if self.size[a] < self.size[b]:
            a, b = b, a
        self.parent[b] = a  # 크기 기준 합치기
        self.size[a] += self.size[b]
        return True

def kruskal(n, edges):
    dsu = DSU(n)
    total, chosen = 0, []
    for w, u, v in sorted(edges):
        if dsu.union(u, v):
            total += w
            chosen.append((u, v, w))
            if len(chosen) == n - 1:
                break
    return total, chosen

edges = [(1, 0, 1), (2, 1, 2), (3, 0, 2), (4, 2, 3), (5, 1, 3)]  # (w, u, v)
print(kruskal(4, edges))  # (7, [(0, 1, 1), (1, 2, 2), (2, 3, 4)])

간선을 (w, u, v)로 저장하면 sorted가 가중치 순으로 정렬합니다. union은 두 끝점이 이미 같은 컴포넌트에 있으면 False를 돌려주는데, 이것이 곧 "사이클이 생기는가" 검사입니다. 고른 간선이 n - 1개보다 적다면 그래프가 연결되어 있지 않은 것이고, 결과는 최소 신장 포레스트입니다.

C++로 옮긴 다익스트라

std::priority_queue는 기본이 최대 힙이므로 greater를 넘겨 최소 힙으로 만듭니다. 합이 넘치지 않도록 long long을 씁니다.

#include <functional>
#include <limits>
#include <queue>
#include <vector>
using namespace std;
using ll = long long;
using State = pair<ll, int>;

vector<ll> dijkstra(const vector<vector<pair<int, ll>>>& adj, int src) {
    vector<ll> dist(adj.size(), numeric_limits<ll>::max());
    priority_queue<State, vector<State>, greater<State>> pq;
    dist[src] = 0;
    pq.push({0, src});
    while (!pq.empty()) {
        auto [d, u] = pq.top();
        pq.pop();
        if (d > dist[u]) continue;
        for (auto [v, w] : adj[u]) {
            if (d + w < dist[v]) {
                dist[v] = d + w;
                pq.push({dist[v], v});
            }
        }
    }
    return dist;
}
// 예제 그래프에서 dijkstra(adj, 0)은 {0, 3, 1, 4, 7}을 돌려줍니다

Java로 옮긴 다익스트라

PriorityQueue는 비교자 기준의 최소 힙입니다. 여기서는 각 항목을 long[] {거리, 정점}으로 둡니다.

import java.util.*;

public class Dijkstra {
    static long[] dijkstra(List<List<int[]>> adj, int src) {
        long[] dist = new long[adj.size()];
        Arrays.fill(dist, Long.MAX_VALUE);
        dist[src] = 0;
        PriorityQueue<long[]> pq = new PriorityQueue<>(Comparator.comparingLong(a -> a[0]));
        pq.add(new long[] {0, src});
        while (!pq.isEmpty()) {
            long[] top = pq.poll();
            long d = top[0];
            int u = (int) top[1];
            if (d > dist[u]) continue;
            for (int[] e : adj.get(u)) {
                int v = e[0];
                if (d + e[1] < dist[v]) {
                    dist[v] = d + e[1];
                    pq.add(new long[] {dist[v], v});
                }
            }
        }
        return dist;
    }
    // 예제 그래프에서 dijkstra(adj, 0)은 [0, 3, 1, 4, 7]을 돌려줍니다
}

TypeScript로 옮긴 다익스트라

JavaScript에는 내장 우선순위 큐가 없으므로 작은 이진 힙을 직접 씁니다. push는 새 항목을 위로 올리고, pop은 마지막 항목을 뿌리로 옮긴 뒤 아래로 내립니다.

type Item = [dist: number, vertex: number];

class MinHeap {
  private a: Item[] = [];
  get size() { return this.a.length; }
  push(item: Item) {
    const a = this.a;
    a.push(item);
    let i = a.length - 1;
    while (i > 0) {
      const p = (i - 1) >> 1;
      if (a[p][0] <= a[i][0]) break;
      [a[p], a[i]] = [a[i], a[p]];
      i = p;
    }
  }
  pop(): Item {
    const a = this.a;
    const top = a[0];
    const last = a.pop()!;
    if (a.length > 0) {
      a[0] = last;
      for (let i = 0; ; ) {
        const l = 2 * i + 1, r = l + 1;
        let m = i;
        if (l < a.length && a[l][0] < a[m][0]) m = l;
        if (r < a.length && a[r][0] < a[m][0]) m = r;
        if (m === i) break;
        [a[m], a[i]] = [a[i], a[m]];
        i = m;
      }
    }
    return top;
  }
}

function dijkstra(adj: [number, number][][], src: number): number[] {
  const dist = new Array<number>(adj.length).fill(Infinity);
  const heap = new MinHeap();
  dist[src] = 0;
  heap.push([0, src]);
  while (heap.size > 0) {
    const [d, u] = heap.pop();
    if (d > dist[u]) continue;
    for (const [v, w] of adj[u]) {
      if (d + w < dist[v]) {
        dist[v] = d + w;
        heap.push([dist[v], v]);
      }
    }
  }
  return dist;
}

console.log(dijkstra([[[1, 4], [2, 1]], [[3, 1]], [[1, 2], [3, 5]], [[4, 3]], []], 0)); // [ 0, 3, 1, 4, 7 ]

정리

• 다익스트라는 힙 반복문입니다. 꺼내고, 낡았으면 건너뛰고, 나가는 간선을 완화하고, 줄어든 거리를 넣습니다.
• 거리뿐 아니라 경로가 필요하면 parent 배열을 둡니다.
• 크루스칼은 "간선을 정렬하고, 아직 연결되지 않았으면 합친다"입니다. 유니온 파인드 덕분에 사이클 검사가 거의 상수 시간입니다.
• C++, Java, TypeScript 버전은 Python 버전과 줄 단위로 대응하며 힙 API만 다릅니다.