백준 1854번: K번째 최단경로 찾기 (Java)

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

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

문제

봄캠프를 마친 김진영 조교는 여러 도시를 돌며 여행을 다닐 계획이다. 그런데 김 조교는, '느림의 미학'을 중요시하는 사람이라 항상 최단경로로만 이동하는 것은 별로 좋아하지 않는다. 하지만 너무 시간이 오래 걸리는 경로도 그리 매력적인 것만은 아니어서, 적당한 타협안인 'k번째 최단경로'를 구하길 원한다. 그를 돕기 위한 프로그램을 작성해 보자.

입력

첫째 줄에 n, m, k가 주어진다. (1 ≤ n ≤ 1000, 0 ≤ m ≤ 2000000, 1 ≤ k ≤ 100) n과 m은 각각 김 조교가 여행을 고려하고 있는 도시들의 개수와, 도시 간에 존재하는 도로의 수이다.

이어지는 m개의 줄에는 각각 도로의 정보를 제공하는 세 개의 정수 a, b, c가 포함되어 있다. 이것은 a번 도시에서 b번 도시로 갈 때는 c의 시간이 걸린다는 의미이다. (1 ≤ a, b ≤ n. 1 ≤ c ≤ 1000)

도시의 번호는 1번부터 n번까지 연속하여 붙어 있으며, 1번 도시는 시작도시이다.

출력

n개의 줄을 출력한다. i번째 줄에 1번 도시에서 i번 도시로 가는 k번째 최단경로의 소요시간을 출력한다.

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

 

---

 

풀이

다익스트라를 돌되, 각 구간의 최단 거리를 dist라는 long 배열에 저장하는 대신, 각 index마다 최단 거리 (최대) K개를 저장하도록 한다. 연산을 하다보면, 각 index마다 저장되어있는 경로의 최댓값만을 들여다볼 일이 많기에, 각 index마다 '(최대) K개의 최단 경로'를 최대힙으로 구현하면 효율적이다.

 

배울 점

솔직히 '풀이 아이디어 떠올리고, 이를 최대 힙으로 구현한다'를 생각해내는게 이 문제의 90%는 될 것이다. 이후 구현은 생각보다 쉬운데, dists[nxt.loc]을 써야할 곳에 pq를 써버리는 바람에 계속 '맞왜틀?'하고 있었다.

생각보다 풀기 어려울 거라고 생각한 플레 4 문제를 풀었다는 생각에 흥분해서 실수를 찾지 못하는 듯하다. 어려운 문제를 푼 거 같아도, 흥분하지 말고 차분히 디버깅하자.

그와 별개로, 다익스트라 문제에 어느정도 도가 튼 듯하다. BFS에서 상태함수 저장하는 느낌이랑 비슷하게 하니까 골드 상위~플레 하위 문제들이 뚝딱이다.

 

코드

import java.io.*;
import java.util.*;
 
public class Main {
    
    static final long INF = Long.MAX_VALUE;
    
    public static void main(String[] args) throws IOException {
        StringBuilder sb = new StringBuilder();
        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());
        
        List<List<Edge>> graph = new ArrayList<>();
        for (int i = 0; i <= N; i++) {
            graph.add(new ArrayList<Edge>());
        }
        
        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());
            graph.get(a).add(new Edge(b, c));
        }
        
        PriorityQueue<Long>[] dists;
        // 각 위치마다 (최대) K개의 최단 경로를 저장할 최대 힙
        dists = (PriorityQueue<Long>[]) new PriorityQueue[N+1];
        for (int i = 0; i <= N; i++) { 
            dists[i] = new PriorityQueue<>(Collections.reverseOrder());
        }
        
        PriorityQueue<Edge> pq = new PriorityQueue<>();
        pq.offer(new Edge(1, 0));
        dists[1].add((long) 0);
        while (!pq.isEmpty()) {
            Edge cur = pq.poll();
            // 해당 위치에 K개의 최소 경로가 있으며 이 중 최댓값보다 큰 경우
            if (dists[cur.loc].size() == K && dists[cur.loc].peek() < cur.dist) continue;
            for (Edge nxt: graph.get(cur.loc)) {
                if (dists[nxt.loc].isEmpty()) {
                    dists[nxt.loc].add(cur.dist + nxt.dist);
                    pq.offer(new Edge(nxt.loc, cur.dist + nxt.dist));
                } else if (dists[nxt.loc].size() < K) {
                    dists[nxt.loc].add(cur.dist + nxt.dist);
                    pq.offer(new Edge(nxt.loc, cur.dist + nxt.dist));
                } else {
                    if (dists[nxt.loc].peek() <= cur.dist + nxt.dist) continue;
                    else {
                        dists[nxt.loc].poll();
                        dists[nxt.loc].add(cur.dist + nxt.dist);
                        pq.offer(new Edge(nxt.loc, cur.dist + nxt.dist));
                    }
                }
            }
        }
        
        for (int i = 1; i <= N; i++) {
            if (dists[i].size() < K) {
                sb.append(-1+"\n");
            }
            else {
                sb.append(dists[i].peek()+"\n");
            }
        }
        System.out.println(sb.toString());
    }
    
    static class Edge implements Comparable<Edge> {
        int loc;
        long dist;
 
        public Edge(int loc, long dist) {
            this.loc = loc;
            this.dist = dist;
        }
 
        public int compareTo(Edge o) {
            return Long.compare(this.dist, o.dist);
        }
    }
}