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

문제 : boj1854

 

  

  다익스트라에 대한 이해를 높히고 으용하는데 매우 도움되는 문제이다. 강추! (플로이드 와샬 이해에는 23258번 밤편지 강추)

 

  다른 문제 풀다가 2번째 최단경로를 찾아야 하는 경우가 있어 찾아보다가 이 문제도 풀게됬다. 다익스트라를 돌릴 때 1부터 n까지 가는 최소 가중치를 알려고 한다고 해보자. 이 때, n번에 처음 도착 했을 때 바로 답을 출력하면 될까? 가중치가 동일한 bfs와 달리 다익스트라는 가중치가 있을 때 사용하기 때문에 그러면 안된다. 다익스트라는 priority queue가 빌 때 까지 다 해봐야 답이 나온다.

 

  그럼 2번째로 작은 가중치를 가지는 경로를 어떻게 구할 수 있을까? 어떠한 정점에 도달한 기록을 2개 유지하면 된다. 물론 이 때에도 마찬가지로 원하는 위치에 2번 도달했을 때 바로 답이 구해지는게 아니다. 모든 위치에서 가중치를 2개 유지하면서 큐가 빌 때 까지 진행해야 한다. 물론 새로 들어온 경로가 현재 위치의 2개 가중치보다 더 작다면 당연히 그걸 넣고 거기서 다시 주변을 큐에 넣어야 한다.

 

  그럼 k번째로 작은 가중치를 가지는 경로도 마찬가지다. 모든 위치에서 해당 지점에 도착한 것 중 최단 가중치 기준으로 k개를 계속 유지하면 된다. 물론 이 때도 현재까지 존재하던 k개의 가중치보다 더 작은 가중치의 값이 들어온다면 그걸로 바꿔주고 거기서 다시 진행해야 한다. 즉 종료되는 조건은 모든 위치 각각에 대해 도달할 방법이 k개 이하였거나, k개가 모두 찼고 더이상 그보다 가중치를 낮출 방법이 없어야 끝난다.

 

  다익스트라를 돌리는 메인 로직은 min heap으로 유지하고, 각 위치에 대해 k개까지만 데이터를 가지는 max heap을 유지하면 k번째 최단경로를 알 수 있다. 어떠한 정점에 해당하는 위치에서 더이상 해당 max heap의 top(k개 중 가장 가중치가 높은 경우)보다 더 가중치가 낮은게 없다면 해당 max heap에 데이터를 추가하지 않는 방식으로 진행하면 된다.

 

 

코드 : github

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class Main {
    class Node implements Comparable<Node> {    // min heap
        int num, c;
        public Node(int num, int t) {
            this.num = num;
            this.c = t;
        }
        @Override
        public int compareTo(Node o) {
            return this.c-o.c;
        }
    }

    private void solve(int n, int k, ArrayList<int[]>[] edges) {
        PriorityQueue<Integer>[] v = new PriorityQueue[n+1];
        for (int i = 1; i <= n; i++) v[i] = new PriorityQueue<>(k, Collections.reverseOrder());

        PriorityQueue<Node> pq = new PriorityQueue<>();
        pq.add(new Node(1, 0));
        v[1].add(0);

        while (!pq.isEmpty()) {
            Node cur = pq.poll();
            if (v[cur.num].size() == k && v[cur.num].peek()<cur.c) continue;
            ArrayList<int[]> edge = edges[cur.num];
            for (int i = 0; i < edge.size(); i++) {
                int[] tmp = edge.get(i);
                int next = tmp[0];
                int nextC = cur.c +tmp[1];
                if (v[next].size() < k) {
                    v[next].add(nextC);
                    pq.add(new Node(next, nextC));
                    continue;
                }
                if (v[next].peek()>nextC) {
                    v[next].poll();
                    v[next].add(nextC);
                    pq.add(new Node(next, nextC));
                    continue;
                }
            }
        }

        StringBuilder sb = new StringBuilder();
        for (int i = 1; i <= n; i++) {
            if (v[i].size() < k)
                sb.append(-1).append('\n');
            else
                sb.append(v[i].peek()).append('\n');
        }
        System.out.print(sb);
    }

    private void solution() throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());
        int k = Integer.parseInt(st.nextToken());
        ArrayList<int[]>[] edges = new ArrayList[n+1];
        for (int i = 1; i <= n ; i++) edges[i] = new ArrayList<>();
        while (m-->0) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            int c = Integer.parseInt(st.nextToken());
            edges[a].add(new int[]{b,c});
        }
        solve(n, k, edges);
    }

    public static void main(String[] args) throws Exception {
        new Main().solution();
    }
}