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

문제

1854번: K번째 최단경로 찾기

풀이

다익스트라, 우선순위 큐 문제이다.

 

처음에는 특정 노드에 도착할 때마다 벡터에 현재 비용을 추가하는 방식으로 접근했다. 하지만 이렇게 하면, 특정 노드에는 k번 도달하지 않고 다익스트라가 종료될 수 있다. 그렇다고, 노드에 도착할 때마다 다익스트라를 돌리는 우선순위 큐에 값을 추가하면 다익스트라가 끝나지 않는다. 

 

다익스트라에 추가적으로 우선순위 큐(내림차순)를 추가하여 풀 수 있다. 1에서 특정 노드까지의 거리를 우선순위 큐로 저장한다.

현재 우선순위 큐에 k개 미만의 수가 채워져있을 때

다익스트라에서 특정 노드에 도달했을 경우, 현재 우선순위 큐의 사이즈와 k를 비교한다. k 미만이라면 현재까지의 거리를 바로 우선순위 큐에 추가한다. 우선순위 큐에 의해 가장 높은 거리가 top에 위치한다. 그래프를 추가적으로 탐색하기 위해 다익스트라를 돌리는 우선순위 큐에도 값을 추가한다. 

현재 우선순위 큐에 k개 이상의 수가 채워져있을 때

k개 이상의 수가 채워져 있다면, 우선순위 큐의 top과 현재까지의 거리를 비교한다. top이 거리보다 크다면, top을 버리고 현재까지의 거리를 우선순위 큐에 추가한다. 그래프를 추가적으로 탐색하기 위해 다익스트라를 돌리는 우선순위 큐에도 값을 추가한다.

이렇게 하면 여러 번 도달할 수 있는 노드가 최소 k번 도달할 때까지 다익스트라가 종료되지 않고, 여러번 도달할 경우에 우선순위 큐의 top이 갱신되지 않을 때까지 다익스트라가 종료되지 않는다.결과적으로, 특정 노드까지의 k번째 거리가 우선순위 큐의 top에 위치하게 된다. 

코드

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

using namespace std;

int k;
vector<vector<pair<int, int>>> edges;
vector<priority_queue<int>> dist_pq;        // 정점까지의 거리를 저장

void dijkstra(int start)
{
    // start에서 start는 0
    dist_pq[start].push(0);

    priority_queue<pair<int, int>> pq;
    pq.push(make_pair(0, start));

    while (!pq.empty())
    {
        int current_node = pq.top().second;
        int current_dist = (-1) * pq.top().first;
        pq.pop();

        for (auto edge : edges[current_node])
        {
            int next_node = edge.second;
            int next_dist = edge.first + current_dist;

            // next_node까지의 거리를 아직 덜 구해놓았다면 dist_pq에 추가한다.
            // 다익스트라를 추가로 진행하기 위해 pq에도 추가한다.
            if (dist_pq[next_node].size() < k)
            {
                dist_pq[next_node].push(next_dist);

                pq.push(make_pair((-1) * next_dist, next_node));
            }
            else 
            {
                // next_node까지의 거리를 이미 k이상 구해놓았다면 현재 가장 높은 비용(top)과 next_dist를 비교하여
                // next_dist가 더 작을 경우에 top을 빼고 next_dist를 추가한다.
                // 다익스트라를 추가로 진행하기 위해 pq에도 추가한다.
                if (dist_pq[next_node].top() > next_dist)
                {
                    dist_pq[next_node].pop();
                    dist_pq[next_node].push(next_dist);

                    pq.push(make_pair((-1) * next_dist, next_node));
                }
            }
        }
    }
}

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

    int n, m;
    cin >> n >> m >> k;

    edges.resize(n + 1);
    dist_pq.resize(n + 1);

    for (int i = 0; i < m; ++i)
    {
        int a, b, c;
        cin >> a >> b >> c;

        edges[a].push_back(make_pair(c, b));
    }

    dijkstra(1);

    for (int i = 1; i <= n; ++i)
    {
        if (dist_pq[i].size() < k)
        {
            cout << "-1" << '\n';
        }
        else
        {
            // dist_pq[i].top()이 i 노드까지의 k번째 비용
            cout << dist_pq[i].top() << '\n';
        }
    }

    return 0;
}