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

1854 : www.acmicpc.net/problem/1854

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

using namespace std;

int n, m, k; // 도시 수, 도시간 간선 수, k번째 최단거리 수
int a, b, c; // a -> b 로 갈때 c만큼 비용
vector<pair<int, int>> vec[1001]; // 인덱스 -> first까지 가는데 second만큼 비용듬
priority_queue<int> city[1001];

int main(){
    // 입력
    scanf("%d %d %d", &n, &m, &k);
    for(int i = 0; i < m; i++){
        scanf("%d %d %d", &a, &b, &c);
        vec[a].push_back({make_pair(b, c)});
    }

    // 다익스트라 알고리즘을 통해 풀이
    priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> q;

    q.push(make_pair(0, 1)); // 우선순위 큐라 비용이 먼저 들어간다
    city[1].push(0);

    while(!q.empty()) {
        int cur, cost;  // 현재지점, 비용
        cur = q.top().second;
        cost = q.top().first;
        q.pop();

        for(int i = 0; i < vec[cur].size(); i++){
            int nextCur, nextCost;
            nextCur = vec[cur][i].first;
            nextCost = vec[cur][i].second;

            if(city[nextCur].size() < k){
                city[nextCur].push(cost + nextCost);
                q.push(make_pair(cost + nextCost, nextCur));
            } else if(city[nextCur].top() > cost + nextCost) {
                city[nextCur].pop();
                city[nextCur].push(cost + nextCost);
                q.push(make_pair(cost + nextCost, nextCur));
            }
        }
    }


    // 출력
    for(int i = 1; i <= n; i++){
        if(city[i].size() < k)
            printf("-1\n");
        else
            printf("%d\n", city[i].top());

    }
}

---

 갓 다익스트라를 배우고 찾아낸 문제지만 생각만 나서 풀이하지 못했는데 찾아보던 중 우선순위 큐를 적절히 사용하는 경우가 있어 풀이를 했습니다.

 

 기존 다익스트라 문제를 풀이할 때 사용하던 배열과 우선순위 큐 외에도 pair로 되어 있는 우선순위 큐를 하나 더 만들어야 합니다. 전역으로 선언된 큐에는 문제의 정답을 뽑아낼 k번째 경우의 수를 저장해 주고 main에 있는 pair 큐를 사용해 다익스트라로 풀어줍니다.

        for(int i = 0; i < vec[cur].size(); i++){
            int nextCur, nextCost;
            nextCur = vec[cur][i].first;
            nextCost = vec[cur][i].second;

            if(city[nextCur].size() < k){
                city[nextCur].push(cost + nextCost);
                q.push(make_pair(cost + nextCost, nextCur));
            } else if(city[nextCur].top() > cost + nextCost) {
                city[nextCur].pop();
                city[nextCur].push(cost + nextCost);
                q.push(make_pair(cost + nextCost, nextCur));
            }
        }

 기존 다익스트라 알고리즘에서는 for문 안에 if 하나로 최우선 경로만 저장하게 되지만 이번 문제는 k 번째를 뽑는게 목적이기 때문에 k개보다 적게 저장되어 있다면 새로 하나를 더 저장하고 만약 k개를 채웠다고 하면 비교 후 큐에 넣어주게 됩니다. 그러면 전역으로 된 큐에는 K개의 값이 있을 것이고 이중 제일 top에 있는게 K번째의 경로가 됩니다.