백준 1162 도로포장 (JAVA)

문제

준영이는 매일 서울에서 포천까지 출퇴근을 한다. 하지만 잠이 많은 준영이는 늦잠을 자 포천에 늦게 도착하기 일쑤다. 돈이 많은 준영이는 고민 끝에 K개의 도로를 포장하여 서울에서 포천까지 가는 시간을 단축하려 한다.
문제는 N개의 도시가 주어지고 그 사이 도로와 이 도로를 통과할 때 걸리는 시간이 주어졌을 때 최소 시간이 걸리도록 하는 K개의 이하의 도로를 포장하는 것이다. 도로는 이미 있는 도로만 포장할 수 있고, 포장하게 되면 도로를 지나는데 걸리는 시간이 0이 된다. 또한 편의상 서울은 1번 도시, 포천은 N번 도시라 하고 1번에서 N번까지 항상 갈 수 있는 데이터만 주어진다.

입력

첫 줄에는 도시의 수 N(1 ≤ N ≤ 10,000)과 도로의 수 M(1 ≤ M ≤ 50,000)과 포장할 도로의 수 K(1 ≤ K ≤ 20)가 공백으로 구분되어 주어진다. M개의 줄에 대해 도로가 연결하는 두 도시와 도로를 통과하는데 걸리는 시간이 주어진다. 도로들은 양방향 도로이며, 걸리는 시간은 1,000,000보다 작거나 같은 자연수이다.

출력

첫 줄에 K개 이하의 도로를 포장하여 얻을 수 있는 최소 시간을 출력한다.

풀이

k개 이하의 도로를 포장하면서 j번 도시까지 가는 최소 시간을 dp[j][k]라고 설정하고, 다익스트라를 돌리면서 최소값을 저장한다. 지금까지 포장한 도로 개수 count에서 1개 더 지웠을 때를 검사하면서(이 경우 다음 edge의 cost는 더하지 않는다) dp를 채워간다.

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

public class Main {
    static int N, M, K;
    static long[][] dp;
    static class Edge {
        int end,count;
        long cost;

        public Edge(int end, long cost, int count) {
            this.end = end;
            this.cost = cost;
            this.count = count;
        }
    }
    static ArrayList<Edge>[] edges;
    static final long MAX = Long.MAX_VALUE;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringBuilder sb = new StringBuilder();
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());
        edges = new ArrayList[N + 1];
        dp = new long[N + 1][K + 1];
        for (int i = 1; i <= N; i++) {
            edges[i] = new ArrayList<>();
            Arrays.fill(dp[i], MAX);
        }
        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int s = Integer.parseInt(st.nextToken());
            int e = Integer.parseInt(st.nextToken());
            int c = Integer.parseInt(st.nextToken());
            edges[s].add(new Edge(e, c, 0));
            edges[e].add(new Edge(s, c, 0));
        }

        dijkstra();
        long min = MAX;
        for (long a : dp[N]) {
            min = Math.min(a, min);
        }
        bw.write(min + "\n");
        bw.flush();
        bw.close();
    }
    static void dijkstra() {
        PriorityQueue<Edge> pq = new PriorityQueue<>(Comparator.comparingLong(o -> o.cost));
        pq.add(new Edge(1, 0, 0));
        dp[1][0] = 0;

        while(!pq.isEmpty()) {
            Edge cur = pq.remove();
            if (cur.cost > dp[cur.end][cur.count]) continue;
            for (Edge next : edges[cur.end]) {
                long nextCost = next.cost + cur.cost;
                if(nextCost < dp[next.end][cur.count]) {
                    dp[next.end][cur.count] = nextCost;
                    pq.add(new Edge(next.end, nextCost, cur.count));
                }
                if (cur.count < K && cur.cost < dp[next.end][cur.count + 1]) {
                    dp[next.end][cur.count + 1] = cur.cost;
                    pq.add(new Edge(next.end, cur.cost, cur.count + 1));
                }
            }
        }
    }
}