[백준 : 1854] K번째 최단경로 찾기 [C/C++] (다익스트라, 우선순위 큐)

 

https://www.acmicpc.net/problem/1854

 

제목에서부터 알 수 있듯이 가중치가 있는 최단경로를 찾는 문제이고, 음수 가중치가 없는 그래프가 주어지기에 다익스트라를 사용하면 된다. 여기서 걸리는 것은 모두가 느꼈듯이 "K번째"라는 점이다.

 

내가 처음으로 바로 떠올린 아이디어는 아래와 같다.

최단경로를 저장할 dist 배열도 배열로 만들고, 이를 정렬하면서 하면 되지 않을까?

근본적으로 맞는 아이디어이다. 그러나 매 탐색마다 정렬을 하면서 하는 것은 의미가 없고, 각각 배열의 사이즈가 K가 넘어가는 경우를 매번 체크하고 저장해야 하기에 시간 복잡도가 커질 위험성이 있다는 것을 깨달았다.

 

적절한 아이디어는 멀지 않은 곳에 존재하고 있었다.

다익스트라에서 자주 사용되는 자료구조인 우선순위 큐를 사용하자!

dist 배열 자체를 우선순위 큐로 만들면 된다. C++의 우선순위 큐는 기본적으로 내림차순 정렬이기에, 나는 이것이 적합하다고 생각했다. 우선순위 큐도 top을 꺼내오는 것이기 때문이다.

 

다익스트라 알고리즘을 그대로 이용하면 된다. 신경써야 될 부분은 각각 노드까지의 최단경로를 저장한 큐의 크기가 k보다 클 경우, 그렇지 않을 경우를 적절하게 구분하여 코드를 작성하는 것이다.

 

#include <iostream>
#include <algorithm>
#include <vector>
#include <queue>

using namespace std;

const int INF = 1e9;

int N,M,K;

int start, dest;

typedef pair<int,int> edge;
int visited[1001];

priority_queue<int> distQ[1001]; // 내림차순
vector<edge> A[1001];
priority_queue<edge,vector<edge>, greater<edge>> q; // 오름차순


void dijkstra(int start){
    q.push(make_pair(0,start)); // 우선순위 큐의 정렬 때문에 순서를 바꿔서 넣는다
    distQ[start].push(0);

    while(!q.empty()){
        int cost = q.top().first;
        int cur = q.top().second;
        q.pop();

        for(int i = 0; i<A[cur].size(); i++){
            int next = A[cur][i].first;
            int n_cost = A[cur][i].second + cost;
			
            // 중요한 부분
            if(distQ[next].size() < K){ // 큐의 사이즈가 K보다 작으면 조건 체크 없이 넣는다
                distQ[next].push(n_cost);
                q.push(make_pair(n_cost,next));
            }
            else if(distQ[next].size() == K && distQ[next].top() > n_cost){
            // 큐의 크기가 K와 같고, 큐의 top(가장 큼)보다 현재 거리가 작으면
                distQ[next].pop(); // top(가장 큼) 제거
                distQ[next].push(n_cost); // 현재 거리 넣음
                q.push(make_pair(n_cost,next));
            }
            // K보다 사이즈가 큰 경우는 존재하지 않는다. 같아질 때 유지만 시키기 때문
        }
    }
}

int main(void){

    cin >> N >> M >> K;

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

    dijkstra(1);

    for(int i = 1; i<=N; i++){
        if(distQ[i].size()==K) cout << distQ[i].top() << "\n"; //사이즈가 같으면(꽉찼으면)
        // 해당 큐 내에서 가장 큰 것이 top일테고, 큐의 사이즈가 K이기 때문에 top이 K번째 최단경로
        else cout << "-1\n";
    }

    return 0;

}

 

잘 이해가 안 가면, 주석과 함께 읽어보길 바란다.