[PS] BOJ 최단경로 - 1753 (C++) | Dijkstra's algorithm

문제 링크 - 최단경로

문제 정리

방향 그래프가 주어졌을 때, 시작 정점 K 에서 다른 모든 정점까지의 최단 경로를 구하는 문제이다. 입력의 범위는 다음과 같다.

• 정점의 개수 $1 \leq V \leq 20{,}000$
• 간선의 개수 $1 \leq E \leq 300{,}000$
• 모든 간선의 가중치는 10 이하의 자연수

전형적인 단일 시작점 최단경로(single-source shortest path) 문제이고, 가중치가 모두 양수이므로 다익스트라 알고리즘을 사용하면 된다.

여기서 먼저 눈여겨 봐야 할 건 V 와 E 의 범위이다. 정점이 20,000 개이므로 인접 행렬을 쓰면 $20{,}000^2 = 4 \times 10^8$ 개의 칸이 필요하고, int 기준으로 1.6GB 라 메모리 제한에 바로 걸린다. 반면 간선은 최대 300,000 개뿐이라 그래프가 굉장히 희소(sparse)하다. 따라서 인접 리스트로 저장해야 한다.

또한 시작 정점에서 도달할 수 없는 정점은 INF 를 출력해야 하고, 같은 정점 쌍 사이에 가중치가 다른 간선이 여러 개 주어질 수 있다는 점도 문제 조건에 명시되어 있다.

다익스트라 알고리즘

다익스트라는 "아직 방문하지 않은 정점 중 시작점에서 가장 가까운 정점"을 하나씩 확정해 나가는 그리디 알고리즘이다. 동작 과정은 다음과 같다.

1. 모든 정점까지의 거리를 INF 로 초기화하고, 시작 정점만 0 으로 둔다.
2. 아직 확정되지 않은 정점 중 거리가 가장 짧은 정점을 하나 고른다.
3. 그 정점을 거쳐 가는 게 더 짧다면 인접한 정점들의 거리를 갱신한다. (relaxation)
4. 모든 정점이 확정될 때까지 2 ~ 3 을 반복한다.

2번에서 고른 정점의 거리가 최단 거리로 확정될 수 있는 이유는 가중치가 모두 양수이기 때문이다. 남은 정점들은 전부 그보다 거리가 멀고, 거기를 한 번 더 거쳐 오면 거리는 더해지기만 하므로 지금보다 짧아질 방법이 없다. 반대로 음수 간선이 있으면 이 전제가 깨지기 때문에 다익스트라를 쓸 수 없고, 벨만-포드를 써야 한다.

왜 우선순위 큐를 쓰는가

2번의 "거리가 가장 짧은 정점 고르기"를 매번 전체 배열을 훑어서 찾으면 한 번에 $O(V)$ 가 걸리고, 이걸 V 번 반복하니 전체 $O(V^2)$ 가 된다. 이 문제에서는 $20{,}000^2 = 4 \times 10^8$ 이라 부담스럽다.

대신 최소 힙(min-heap)을 쓰면 가장 짧은 정점을 $O(\log V)$ 에 꺼낼 수 있다. 간선을 갱신할 때마다 힙에 넣으므로 전체 시간 복잡도는 다음과 같다.

넉넉하게 통과한다.

해결 전략

C++ 의 priority_queue 는 기본이 최대 힙이므로, 최소 힙으로 바꿔줘야 한다.

priority_queue<pair<int, int>, vector<pair<int, int>>, greater<>> pq;

큐에는 {거리, 정점} 순서로 넣는다. pair 는 first 를 먼저 비교하기 때문에 이 순서로 넣어야 거리 기준으로 정렬된다. (정점 번호를 앞에 넣으면 엉뚱하게 정점 번호순으로 정렬된다)

그리고 한 가지 더 중요한 처리가 있다. 거리가 갱신될 때마다 큐에 새로 넣기 때문에, 같은 정점이 서로 다른 거리를 가진 채로 큐에 여러 번 들어갈 수 있다. 이미 더 짧은 거리로 처리가 끝난 정점이 나중에 긴 거리로 다시 나오는 것인데, 이런 항목은 버려야 한다.

if(cost > dist[cur]) continue;

이 한 줄이 별도의 visited 배열 역할을 대신한다. dist[cur] 는 항상 지금까지 찾은 최솟값이므로, 큐에서 꺼낸 거리가 그보다 크다면 이미 한물간(stale) 데이터라는 뜻이다.

예시

문제의 예제 입력을 따라가 보자.

5 6
1
5 1 1
1 2 2
1 3 3
2 3 4
2 4 5
3 4 6

정점 1 에서 시작하므로 dist[1] = 0 이고 나머지는 INF 이다.

정점 2 에서 정점 3 으로 가는 경로는 $2 + 4 = 6$ 이지만 이미 dist[3] = 3 이므로 갱신되지 않는다. 정점 5 는 5 → 1 간선만 있고 들어오는 간선이 없어서 끝까지 INF 로 남는다.

0
2
3
7
INF

문제 풀이

구현하면서 고려한 내용들을 정리해 보았다.

• 위에서 이야기한 대로 메모리 때문에 인접 행렬 대신 인접 리스트(vector<pair<int, int>> graph[])를 사용했다. 중복 간선도 그냥 리스트에 같이 담아두면 되는데, 어차피 relaxation 과정에서 더 짧은 쪽만 살아남기 때문에 따로 걸러줄 필요가 없다.
• INF 는 1e9 로 잡았다. 최대 거리가 $20{,}000 \times 10 = 200{,}000$ 이라 이보다 훨씬 크고, cost + weight 를 계산해도 int 범위를 넘지 않는다.
• 큐에서 꺼낸 항목이 이미 처리된 것인지 cost > dist[cur] 로 검사해 걸렀다.
• auto [cost, cur] = pq.top(); 같은 구조적 바인딩은 C++17 부터 사용 가능하다.

코드

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
using namespace std;

const int INF = 1e9;
int V, E, K;
vector<pair<int, int>> graph[20001];
int dist[20001];

void dijkstra(int start)
{
    priority_queue<pair<int, int>, vector<pair<int, int>>, greater<>> pq;

    fill(dist + 1, dist + V + 1, INF);
    dist[start] = 0;
    pq.push({0, start});

    while(!pq.empty())
    {
        auto [cost, cur] = pq.top(); pq.pop();

        // 이미 더 짧은 경로로 처리된 정점이면 건너뛴다
        if(cost > dist[cur]) continue;

        for(auto [next, weight] : graph[cur])
        {
            if(cost + weight < dist[next])
            {
                dist[next] = cost + weight;
                pq.push({dist[next], next});
            }
        }
    }
}

void solve()
{
    cin >> V >> E >> K;

    for(int i = 0; i < E; i++)
    {
        int u, v, w; cin >> u >> v >> w;
        graph[u].push_back({v, w});
    }

    dijkstra(K);

    for(int i = 1; i <= V; i++)
    {
        if(dist[i] == INF) cout << "INF" << '\n';
        else cout << dist[i] << '\n';
    }
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);

    solve();

    return 0;
}

정리

• 가중치가 양수인 단일 시작점 최단경로 → 다익스트라
• 정점은 많고 간선은 상대적으로 적으면 → 인접 리스트 + 우선순위 큐
• priority_queue 는 기본이 최대 힙이므로 greater<> 로 최소 힙을 만들고, {거리, 정점} 순서로 넣기
• 큐에 중복으로 들어간 항목은 cost > dist[cur] 로 걸러내기

시간 복잡도는 $O(E \log V)$, 공간 복잡도는 $O(V + E)$ 이다. 음수 가중치가 섞여 있다면 이 풀이는 쓸 수 없으니, 그때는 벨만-포드나 SPFA 를 사용해야 한다.