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

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

 

문제

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

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

 

 

풀이

다익스트라 알고리즘과 우선순위 큐(최대 힙)를 사용해 풀 수 있는 문제다. 시작 지점이 1로 고정되어 다른 모든 정점까지의 거리를 구하는 문제이므로 다익스트라 알고리즘을 쓸 수 있다. 원래 다익스트라 알고리즘은 최단 거리를 저장하는 dist 배열의 값을 비교해 갱신해 나간다. 하지만 이 문제에서는 K번째 거리를 구해야 하므로 최솟값만 저장하는 배열 대신 우선순위 큐의 배열이 필요하다.

반응형

따라서 우선순위 큐(최대 힙) 배열인 h(heap)를 만들어주고 다익스트라 알고리즘을 수행하며 어느 정점에 도착했을 때 그 정점의 최대 힙의 크기가 K보다 작다면 거리를 추가해 주고, K 이상(실제로는 K로 유지)인 경우에는 거리를 top 원소와 비교해 top 원소보다 작은 경우에만 pop 하고 거리를 push 해주어 top 원소가 K번째 최단 경로로 유지될 수 있도록 한다.

 

다익스트라 알고리즘이 끝난 후, 각 정점의 최대 힙에 대해 크기가 K 미만인 경우 -1을, K인 경우 top 원소를 출력한다.

 

 

코드

#include <bits/stdc++.h>
using namespace std;
#define pii pair<int, int>

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);
  int V, E, K; // 정점, 간선, K번째
  vector<pii> graph[1001];
  cin >> V >> E >> K;
  while (E--) {
    int u, v, w;
    cin >> u >> v >> w;
    graph[u].push_back({v, w});
  }
  priority_queue<int> h[1001]; // K개의 수를 저장할 최대힙 배열
  // 다익스트라 알고리즘
  priority_queue<pii> pq;
  h[1].push(0); // 시작 지점 비용 0 추가
  pq.push({0, 1});
  while (!pq.empty()) {
    int cost = -pq.top().first;
    int curr = pq.top().second;
    pq.pop();
    for (pii p : graph[curr]) {
      int next = p.first, d = cost + p.second;
      if (h[next].size() < K) {
     	// 힙의 크기가 K보다 작다면 그냥 추가
        h[next].push(d);
        pq.push({-d, next});
      } else if (h[next].top() > d) {
      	// 힙의 크기가 K이면서
        // top 원소(현재 K번째 비용)보다 현재 비용이 더 작다면
        // top 원소를 pop하고 현재 비용을 push
        h[next].pop();
        h[next].push(d);
        pq.push({-d, next});
      }
    }
  }
  // 각 정점에 대해 최대힙의 크기가 K보다 작다면 -1, 그렇지 않으면 최대힙의 top을 출력
  for (int i = 1; i <= V; i++)
    cout << (h[i].size() < K ? -1 : h[i].top()) << '\n';
  return 0;
}