자바 PriorityQueue로 다익스트라 최단경로 구현

코딩테스트 그래프 문제를 풀다 보면 다익스트라가 거의 매번 등장한다. 그런데 개념은 아는데 막상 제출하면 시간초과(TLE)로 미끄러지는 경우가 많다. 원인을 따라가 보면 대부분 PriorityQueue를 쓰는 방식이 아니라 이미 확정된 노드를 큐에서 걸러내지 않는 것에 있다. 이 글은 그 지점을 코드로 짚는다.

결론부터 말하자면, 우선순위큐 기반 다익스트라의 핵심은 두 가지다. 첫째, 그래프를 인접리스트로 표현해 간선 순회를 O(E)로 줄인다. 둘째, 큐에서 꺼낸 거리값이 이미 갱신된 최단거리보다 크면 즉시 버린다. 이 두 가지만 지키면 정점 V, 간선 E에 대해 O(E log V)로 동작한다.

1. 다익스트라와 우선순위큐가 하는 일

다익스트라는 음의 간선이 없는 그래프에서 한 출발점부터 모든 정점까지의 최단거리를 구하는 알고리즘이다. 매 단계에서 '아직 확정되지 않은 정점 중 거리가 가장 짧은 것'을 골라 확정하고, 그 정점을 거쳐 이웃으로 가는 거리를 갱신한다.

'가장 짧은 것을 고르는' 연산을 매번 배열 전체에서 선형 탐색하면 O(V^2)이 된다. 여기서 PriorityQueue가 등장한다. 자바의 PriorityQueue는 이진 힙 기반 최소/최대 힙으로, offer와 poll이 O(log n), peek이 O(1)이다. 최단거리 후보를 힙에 넣어두면 매번 최소값을 log 시간에 꺼낼 수 있다.

기본 정렬은 자연순서(오름차순)이고, 원소가 Comparable이 아니거나 다른 기준이 필요하면 생성자에 Comparator를 넘긴다. 다익스트라에서는 '거리가 작은 것 우선'이므로 거리 기준 오름차순 비교자를 쓴다.

2. 인접리스트로 그래프 표현하기

정점 수가 크면 인접행렬 int[V][V]는 메모리도 O(V^2)이고 간선 순회도 느리다. 간선이 희소한 코딩테스트 그래프에서는 인접리스트가 맞다. 정점마다 나가는 간선 목록만 담는다.

// 간선: 목적지 to, 가중치 weight
static class Edge {
    int to;
    int weight;
    Edge(int to, int weight) {
        this.to = to;
        this.weight = weight;
    }
}

// 정점 개수 n, 1-indexed 기준
List<List<Edge>> graph = new ArrayList<>();
for (int i = 0; i <= n; i++) graph.add(new ArrayList<>());

// u -> v 방향, 가중치 w 간선 추가
graph.get(u).add(new Edge(v, w));

무방향 그래프라면 graph.get(v).add(new Edge(u, w))도 함께 넣는다. 방향을 한쪽만 넣고 무방향 문제를 푸는 실수가 흔하다.

3. PriorityQueue 다익스트라 전체 코드

큐에는 [정점 번호, 그 정점까지의 현재 거리]를 int[]로 담고, 거리 기준으로 비교한다. 배열 대신 int[]를 쓰면 오토박싱 부담이 줄어 실전에서 유리하다.

import java.util.*;

public class Dijkstra {
    static final int INF = Integer.MAX_VALUE;

    static int[] dijkstra(List<List<Edge>> graph, int start, int n) {
        int[] dist = new int[n + 1];
        Arrays.fill(dist, INF);
        dist[start] = 0;

        // int[]{정점, 거리}, 거리 오름차순
        PriorityQueue<int[]> pq =
            new PriorityQueue<>((a, b) -> a[1] - b[1]);
        pq.offer(new int[]{start, 0});

        while (!pq.isEmpty()) {
            int[] cur = pq.poll();
            int node = cur[0];
            int d = cur[1];

            // 핵심: 이미 더 짧은 경로로 확정된 노드면 버린다
            if (d > dist[node]) continue;

            for (Edge e : graph.get(node)) {
                int next = e.to;
                int cost = d + e.weight;
                if (cost < dist[next]) {
                    dist[next] = cost;
                    pq.offer(new int[]{next, cost});
                }
            }
        }
        return dist;
    }
}

비교자를 a[1] - b[1]로 쓴 이유는 거리값이 int 범위 안이라 오버플로가 없기 때문이다. 거리 합이 int를 넘길 수 있는 문제라면 뒤에서 다시 다룬다.

4. 실행 결과

정점 5개, 시작점 1번인 예제로 돌려본다.

int n = 5;
List<List<Edge>> graph = new ArrayList<>();
for (int i = 0; i <= n; i++) graph.add(new ArrayList<>());
graph.get(1).add(new Edge(2, 2));
graph.get(1).add(new Edge(3, 5));
graph.get(2).add(new Edge(3, 1));
graph.get(2).add(new Edge(4, 6));
graph.get(3).add(new Edge(4, 2));
graph.get(4).add(new Edge(5, 1));

int[] dist = dijkstra(graph, 1, n);
for (int i = 1; i <= n; i++)
    System.out.println(i + ": " + (dist[i] == INF ? "INF" : dist[i]));

1: 0
2: 2
3: 3
4: 5
5: 6

1번에서 3번으로 갈 때 직접 간선(5)보다 2번을 경유하는 경로(2+1=3)가 짧게 잡힌다. 4번도 1-2-3-4(2+1+2=5)로 갱신된다. 갱신이 제대로 전파되는지 이 예제로 먼저 확인하고 넘어가면 디버깅이 편하다.

5. 방문 배열 최적화: dist 비교로 대체한다

BFS 습관대로 boolean[] visited를 따로 두는 코드를 자주 본다. 다익스트라에서는 별도 방문 배열이 필수가 아니다. dist 배열 자체가 '지금까지 찾은 최단거리'를 들고 있으므로, 큐에서 꺼낸 거리 d가 dist[node]보다 크면 그 항목은 낡은 정보다. 바로 continue로 건너뛰면 된다.

if (d > dist[node]) continue;   // 이 한 줄이 방문 처리 겸 최적화

이 조건이 하는 일은 두 가지다. 이미 확정된 정점을 다시 펼치지 않게 막고, 뒤에서 설명할 중복 삽입 항목을 싼값에 폐기한다. visited를 굳이 쓰겠다면 poll 직후 if (visited[node]) continue; visited[node] = true; 형태로 두는데, 동작은 같고 코드만 늘어난다.

6. 자주 틀리는 지점 1: 큐 중복 삽입과 지연 삭제

우선순위큐 다익스트라는 같은 정점을 큐에 여러 번 넣는다. 더 짧은 경로를 발견할 때마다 offer하기 때문이다. 자바 PriorityQueue에는 힙 안에서 특정 원소의 우선순위를 O(log n)에 낮추는 decrease-key 연산이 없다. 그래서 값을 고치는 대신 새 항목을 밀어넣고, 꺼낼 때 낡은 항목을 버리는 '지연 삭제' 방식을 쓴다.

여기서 if (d > dist[node]) continue;를 빼먹으면 낡은 항목까지 전부 이웃 순회를 돌게 된다. 큐 크기가 불필요하게 커지고, 최악의 경우 재갱신이 연쇄되면서 시간초과로 이어진다. 중복 삽입 자체는 정상이고, 꺼낼 때 거르는 코드가 없는 것이 문제다.

7. 자주 틀리는 지점 2: 시간초과(TLE)의 실제 원인

제출 후 TLE가 뜨면 대개 아래 중 하나다.

• 인접행렬 사용: 정점이 수만 개면 int[V][V] 순회만으로 터진다. 인접리스트로 바꾼다.
• 낡은 항목 미폐기: 위 6절의 continue 누락. 가장 흔하다.
• Scanner로 대량 입력: 간선이 수십만 줄이면 Scanner 파싱이 병목이다. BufferedReader와 StringTokenizer로 바꾼다.
• System.out.println 반복 출력: 결과를 한 줄씩 출력하면 느리다. StringBuilder에 모아 한 번에 출력한다.

BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringTokenizer st = new StringTokenizer(br.readLine());
int n = Integer.parseInt(st.nextToken());
int m = Integer.parseInt(st.nextToken());

알고리즘이 O(E log V)로 맞는데도 TLE라면 알고리즘보다 입출력을 먼저 의심한다.

8. 자주 틀리는 지점 3: 거리 오버플로와 비교자

거리 초기값을 Integer.MAX_VALUE로 두고 간선 가중치가 크면, d + e.weight 계산에서 int 오버플로가 날 수 있다. 갱신 조건 cost < dist[next]가 있어 INF에 더하는 상황은 대부분 걸러지지만, 경로 합 자체가 int 범위를 넘는 문제라면 dist와 거리 계산을 long으로 바꾼다.

이때 비교자도 주의한다. long에서는 a[1] - b[1]식 뺄셈 비교가 오버플로로 부호가 뒤집힐 수 있다. Long.compare(a[1], b[1])를 쓴다.

PriorityQueue<long[]> pq =
    new PriorityQueue<>((a, b) -> Long.compare(a[1], b[1]));

9. 정리

우선순위큐 다익스트라를 안정적으로 통과시키려면 인접리스트로 그래프를 표현하고, 큐에서 꺼낸 항목이 낡았는지 d > dist[node]로 검사해 버리며, 대량 입출력은 BufferedReader와 StringBuilder로 처리한다. 별도 방문 배열 없이 dist 비교만으로 방문 처리가 되고, 거리 합이 커지는 문제에서는 long과 Long.compare로 오버플로를 막는다.