[백준] 1854번 K번째 최단경로 찾기

[적용한 알고리즘]

다익스트라

 

[아이디어]

dist 배열을 최단경로를 저장하는 곳이 아닌,

k 개의 '최단경로들' 을 저장하는 용도로 사용한다.

(Heap 을 이용해서)

 

따라서 dist[i].top() 은 출발점에서 i 로 가는 k 번째 최단경로이다.

그래서 기존의 다익스트라에서 최단경로를 저장하는 부분을 살짝 손봐줘야하는데,

 

dist[nxt].size() < K 이거나 dist[nxt].top() > (현재 정점까지 k 번째 최단경로) + cost 라면,

dist[nxt] 에 ((현재 정점까지 k 번째 최단경로) + cost) push,

원래의 우선순위 큐에는 {(현재 정점까지 k 번째 최단경로) + cost, 다음 정점) 을 push 해주면 되겠다.

그리고 만약 size() 가 k 였다면, 하나를 당연히 pop 해줘야한다.

 

끝까지 돌았을 때 경로의 수가 k 개가 안된다면, 그 경로가 없는 것이므로 -1

있다면 top 을 출력하면 된다.

 

[교훈]

개 어 렵..

 

<코드>

#include <cstdio>
#include <queue>
using namespace std;
typedef pair<int, int> P;
const int MAX_V = 10001;
int N, M, K;

int main(){
    scanf("%d %d %d", &N, &M, &K);
    vector<pair<int, int>> adj[MAX_V];

    for (int i = 0; i < M; i++){
        int from, to, cost;
        scanf("%d %d %d", &from, &to, &cost);
        adj[from].push_back({to, cost});
    }

    priority_queue<int> dist[MAX_V];
    priority_queue<P, vector<P>, greater<P>> PQ;
    dist[1].push(0);
    PQ.push({0, 1});

    while (!PQ.empty()){
        auto p = PQ.top();
        PQ.pop();

        int cur = p.second;
        int dist_cur = p.first;
        
        for (auto q : adj[cur]){
            int nxt = q.first;
            int cost = q.second;

            if (dist[nxt].size() < K || dist[nxt].top() > dist_cur + cost){
                if (dist[nxt].size() == K){
                    dist[nxt].pop();
                }
                dist[nxt].push(dist_cur + cost);
                PQ.push({dist_cur + cost, nxt});
            }
        }
    }

    for (int i = 1; i <= N; i++){
        if (dist[i].size() != K){
            printf("-1\n");
        }
        else {
            printf("%d\n", dist[i].top());
        }
    }
    return 0;
}