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

문제

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

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

 

설명

i번 도시에서 i번 도시로 가는 최단경로는 0이지만, 일반적인 k번째 최단경로는 0이 아닐 수 있음에 유의한다.
또, k번째 최단경로가 존재하지 않으면 -1을 출력한다. 최단경로에 같은 정점이 여러 번 포함되어도 된다.

도로의 정보를 저장하기 위한 Road 클래스를 만들고 리스트를 만들어 연결관계를 저장한다.

이때, 소요 시간이 작은 순으로 관리하기 위해 Comparable 인터페이스를 이용한다.

k번째 최단경로를 구해야 하므로 각 도시마다 우선순위 큐를 만들어 k개까지의 경로를 저장해 관리한다.

맨 앞에 가장 긴 경로가 오도록 Collections.reverseOrder()를 이용한다. 

class Road implements Comparable<Road> {
    int dest;
    int time;

    public Road(int dest, int time) {
        this.dest = dest;
        this.time = time;
    }

    @Override
    public int compareTo(Road o) {
        return this.time - o.time;
    }
}

List<List<Road>> path = new ArrayList<>();
List<PriorityQueue<Integer>> pathPQ = new ArrayList<>();
for (int i = 0; i <= N; i++) {
    path.add(new ArrayList<>());
    pathPQ.add(new PriorityQueue<>(Collections.reverseOrder()));    // 역순 - 내림차순
}

음의 가중치가 없는 그래프이므로 다익스트라 알고리즘을 사용해 최단경로를 구한다.

방문하는 정점을 관리하기 위해 우선순위 큐를 하나 만든다.

시작 정점과 소요 시간을 각각 1과 0으로 하여 우선순위 큐에 삽입한다.

위에서 만든 pathPQ에서도 1번 도시에 해당하는 우선순위 큐에 0을 삽입한다.

PriorityQueue<Road> pq = new PriorityQueue<>();
pq.offer(new Road(1, 0));
pathPQ.get(1).offer(0);

우선순위 큐가 비어있을 때까지 최단 경로를 갱신한다.

먼저 우선순위 큐에서 하나를 빼내어 해당 dest 도시와 연결된 도로들을 조사한다.

이때, nextRD 도시에 해당하는 우선순위 큐에 k개의 경로가 있는 경우와 그렇지 않은 경우로 나뉜다.

k개보다 작은 경로가 저장되어 있다면 pathPQ에 해당 경로를 저장하고 pq에 해당 도로를 삽입한다.

k개의 경로가 저장되어 있다면 현재 우선순위 큐 앞에는 가장 긴 경로가 저장되어 있으므로 peek()과 해당 경로와 비교하여 크기가 더 작은 경로를 저장한다.

모든 과정이 끝난 후에 k번째 최단 경로가 없다면 -1을 출력하고 존재한다면 각 도시에 해당하는 우선순위 큐의 peek()을 출력한다.

void findMinLength() {
    while (!pq.isEmpty()) {
        Road road = pq.poll();
        int RD = road.dest;
        int RT = road.time;

        int len = path.get(RD).size();

        // 인접 리스트로 비교
        for (int i = 0; i < len; i++) {
            int nextRD = path.get(RD).get(i).dest;
            int nextRT = path.get(RD).get(i).time;

            // pathPQ[끝점과 인접한 점]의 사이즈 < K (K번째 최단경로의 소요시간)
            if (pathPQ.get(nextRD).size() < K) {
                pathPQ.get(nextRD).offer(nextRT + RT);
                pq.offer(new Road(nextRD, nextRT + RT));
            }
            else {
                // K번째 최단 경로는 항상 pathPQ.get(nextRD)의 top으로 유지되어 있다
                if (pathPQ.get(nextRD).peek() > nextRT + RT) {
                    // pathPQ의 제일 윗값 (max 값을 제거)
                    pathPQ.get(nextRD).poll();

                    // pathPQ에 신규 값을 넣어준다
                    pathPQ.get(nextRD).offer(nextRT + RT);
                    pq.offer(new Road(nextRD, nextRT + RT));
                }
            }
        }
    }

    for (int i = 1; i <= N; i++) {
        if (pathPQ.get(i).size() != K)
            System.out.println(-1);
        else
            System.out.println(pathPQ.get(i).peek());
    }
}

 

전체 코드

import java.io.*;
import java.util.*;

public class Main {

    static int N, M, K;
    static List<List<Road>> path;
    static List<PriorityQueue<Integer>> pathPQ;
    static PriorityQueue<Road> pq;

    static class Road implements Comparable<Road> {
        int dest;
        int time;

        public Road(int dest, int time) {
            this.dest = dest;
            this.time = time;
        }

        @Override
        public int compareTo(Road o) {
            return this.time - o.time;
        }
    }

    static void findMinLength() {
        while (!pq.isEmpty()) {
            Road road = pq.poll();
            int RD = road.dest;
            int RT = road.time;

            int len = path.get(RD).size();

            // 인접 리스트로 비교
            for (int i = 0; i < len; i++) {
                int nextRD = path.get(RD).get(i).dest;
                int nextRT = path.get(RD).get(i).time;

                // pathPQ[끝점과 인접한 점]의 사이즈 < K (K번째 최단경로의 소요시간)
                if (pathPQ.get(nextRD).size() < K) {
                    pathPQ.get(nextRD).offer(nextRT + RT);
                    pq.offer(new Road(nextRD, nextRT + RT));
                }
                else {
                    // K번째 최단 경로는 항상 pathPQ.get(nextRD)의 top으로 유지되어 있다
                    if (pathPQ.get(nextRD).peek() > nextRT + RT) {
                        // pathPQ의 제일 윗값 (max 값을 제거)
                        pathPQ.get(nextRD).poll();

                        // pathPQ에 신규 값을 넣어준다
                        pathPQ.get(nextRD).offer(nextRT + RT);
                        pq.offer(new Road(nextRD, nextRT + RT));
                    }
                }
            }
        }

        for (int i = 1; i <= N; i++) {
            if (pathPQ.get(i).size() != K)
                System.out.println(-1);
            else
                System.out.println(pathPQ.get(i).peek());
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());

        path = new ArrayList<>();
        pathPQ = new ArrayList<>();
        for (int i = 0; i <= N; i++) {
            path.add(new ArrayList<>());
            pathPQ.add(new PriorityQueue<>(Collections.reverseOrder()));    // 역순 - 내림차순
        }

        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int A = Integer.parseInt(st.nextToken());
            int B = Integer.parseInt(st.nextToken());
            int C = Integer.parseInt(st.nextToken());

            path.get(A).add(new Road(B, C));
        }

        pq = new PriorityQueue<>();
        pq.offer(new Road(1, 0));
        pathPQ.get(1).offer(0);

        findMinLength();
    }
}