[백준 P1854] K번째 최단경로 찾기 Java

P1854 K번째 최단경로 찾기

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번 도시는 시작도시이다.

5 10 2
1 2 2
1 3 7
1 4 5
1 5 6
2 4 2
2 3 4
3 4 6
3 5 8
5 2 4
5 4 1

출력값

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

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

-1
10
7
5
14

알고리즘

일반적으로 출발지에서 모든 정점의 최단 거리를 구할 때에는 다익스트라나 벨만 포드 알고리즘을 활용한다. 간선의 가중치(c)는 입력 값의 범위 조건에 따라서 항상 양수이므로, 다익스트라 알고리즘을 사용하여 풀이한다.

 

K번째 최단 경로는 시작점에서 종점으로 갈 수 있는 모든 경로 중에서 K번째로 작은 가중치를 가지는 경로를 뜻한다. 문제에서 원하는 것은 이 K번째 최단 경로를 구하는 것이 아닌, K번째로 작은 경로의 가중치를 구하는 것이므로 따로 경로를 저장할 필요는 없이, 모든 정점에 대해 K번째로 작은 값만 뽑아낼 수 있으면 된다.

 

일반적인 최단 경로 문제는 한 싸이클에 하나의 최단 경로를 확정하면서 이미 방문한 노드는 다시 방문하지 않도록 방문 처리를 해준다. 하지만, 이 문제에서는 최단 경로에 같은 정점이 여러 번 포함될 수 있으므로 코드의 일부 수정이 필요하다. 다익스트라의 일반적인 풀이는 여기서 언급하지 않겠다. 우선순위 큐를 사용했으며 만약 모르는 개념이라면 글을 읽기 전에 이 포스팅을 참고하면 될 것 같다.

[백준 P1753] 최단경로 Java

 

2개의 자료구조를 사용한다.

1. 이후에 탐색할 경로들이 등록되어 있는 간선 우선순위 큐가 필요하다. 다익스트라 알고리즘을 수행할 때 간선을 저장하는 용도로 사용된다. 글에서는 pq로 표현하겠다.
2. 각 정점에 대해 특정 시점까지 탐색한 경로의 값들을 저장할 자료구조가 필요하다. 모든 간선을 탐색한 이후에 각 정점에 대해 K번째 최단 경로의 값을 얻기 위한 용도로 사용된다. 이는 우선순위 큐의 배열로 구현했다. 글에서는 dest로 표현하겠다.

 

다익스트라 알고리즘을 수행하면서 pq와 dest를 업데이트하는 과정에서 매 간선에 대해 두 가지 조건을 확인해야 한다.

1. 출발지에서 목적지까지 구해둔 최단 경로의 수 (dest[목적지].size())가 K개 미만인가?
2. 출발지에서 목적지까지 구해둔 최단 경로의 수 (dest[목적지].size())가 K개 이상이라면 구해둔 K개의 최단 경로 중 가장 짧은 경로(dest[목적지].peek()가 새롭게 업데이트 되는 경로보다 긴 값인가?

 

조건 1과 2를 더 빠르게 수행하기 위해서 2번 자료구조는 우선순위 큐의 배열로 구현했다. 조건 1은 PriorityQueue.size() 메서드를 사용하고, 조건 2는 PriorityQueue.peek() 메서드를 사용한다.

 

일반적인 다익스트라에서는 다음 경로를 추가할 때 지금까지 찾은 경로가 저장된 경로보다 더 작을 때만 값을 업데이트 하고 이후에 탐색할 경로로 등록한다. 그 부분을 위의 조건에 맞게 지금까지 찾은 경로가 K개보다 작거나, K개 이상이라면 K개 중에서 가장 큰 값과 비교해서 작을 때만 dest[목적지]에 값을 추가하고 pq에 간선을 추가하도록 바꾸었다.

for (Edge next: graph[current.to]) {
    if (dest[next.to].size() < K || dest[next.to].peek() > current.weight + next.weight) {
        if (dest[next.to].size() == K) dest[next.to].poll();
        dest[next.to].add(current.weight + next.weight);
        pq.add(new Edge(next.to, current.weight + next.weight));
    }
}

예시를 통해 알고리즘이 어떻게 동작할지 생각해보자. 아래 그래프에서 2번째 최단 경로를 구한다고 가정하자. 1번 정점을 시작점으로 보겠다.

시작점에서 시작점까지 거리를 0으로 설정하여 dest[0]에 0을 넣어주고 pq에 (1, 1, 0) 간선을 넣는다.

 

 

최소 가중치의 간선을 pq에서 하나 꺼낸다. (1, 1, 0)에서 갈 수 있는 모든 간선에 대해 조건 1, 2를 체크한다. 만약 조건 1, 2를 만족한다면 새로운 간선을 pq에 추가하면서 dest[목적지]에 값을 추가한다.

현재 (1, 1, 0) 간선을 통해 탐색할 수 있는 새로운 경로는 (1, 2, 2), (1, 3, 5)이다. 지금까지 2와 3으로 갈 수 있는 경로의 수는 0으로 K 보다 작다. 따라서 두 간선 모두 pq에 넣어주고, dest[2]과 dest[3]에 각 값을 추가해준다.

 

최소 가중치의 간선을 pq에서 하나 꺼낸다. (1, 2, 2)에서 갈 수 있는 모든 간선에 대해 조건 1, 2를 체크한다. 현재 (1, 2, 2) 간선을 통해 탐색할 수 있는 새로운 경로는 (2, 1, 4), (2, 3, 3)이다. 지금까지 1과 3으로 갈 수 있는 경로는 각각 0(1개)과 5(1개)로 K개보다 작다. 따라서 두 간선 모두 pq에 넣어주고, dest[1]과 dest[3]에 각 값을 추가해준다.

 

최소 가중치의 간선을 pq에서 하나 꺼낸다. (2, 3, 3)에서 갈 수 있는 모든 간선에 대해 조건 1, 2를 체크한다. 현재 (2, 3, 3) 간선을 통해 탐색할 수 있는 새로운 경로는 (3, 2, 4), (3, 1, 8)이다. 지금까지 1과 2으로 갈 수 있는 경로는 각각 0, 4(2개)와 2(1개)개 이다. 2번 정점으로 가는 경우의 수는 아직 1개로 K보다 작기 때문에 (3, 2, 4) 간선은 pq에 넣어주고 dest[2]에 값을 추가해준다.

현재 1번 정점으로 가는 경우의 수는 K개와 같다. 따라서 현재 1번으로 가는데 필요한 가중치 8과 지금까지 1번으로 가기 위해 저장한 가중치 값 중 가장 큰 값(dest[1].peek())을 비교하여 더 작은 값을 저장한다. 4와 8 중에 4가 더 작으므로 기존 값이 더 작다. 따라서 (3, 1, 8)은 pq에 추가하지 않고 dest[1]도 업데이트 하지 않는다.

 

최소 가중치의 간선을 pq에서 하나 꺼낸다. (2, 1, 4)에서 갈 수 있는 모든 간선에 대해 조건 1, 2를 체크한다. 현재 (2, 1, 4) 간선을 통해 탐색할 수 있는 새로운 경로는 (1, 2, 6), (1, 3, 9)이다. 지금까지 2과 3으로 갈 수 있는 경로는 각각 2, 4(2개)과 3, 5(2개)로 K개이다.

현재 2번 정점으로 가는 경우의 수는 K개와 같다. 따라서 현재 2번으로 가는데 필요한 가중치 6과 지금까지 2번으로 가기 위해 저장한 가중치 값 중 가장 큰 값(dest[2].peek())을 비교하여 더 작은 값을 저장한다. 6과 6은 같으므로 기존 값을 변경하지 않는다. 따라서 (1, 2, 6)은 pq에 추가하지 않고 dest[2]도 업데이트 하지 않는다.

현재 3번 정점으로 가는 경우의 수는 K개와 같다. 따라서 현재 3번으로 가는데 필요한 가중치 9과 지금까지 3번으로 가기 위해 저장한 가중치 값 중 가장 큰 값(dest[3].peek())을 비교하여 더 작은 값을 저장한다. 5과 9 중에 기존 값인 5가 더 작으므로 기존 값을 변경하지 않는다. 따라서 (1, 3, 9)은 pq에 추가하지 않고 dest[3]도 업데이트 하지 않는다.

 

 

위를 반복하면서 pq에 남아있는 모든 간선을 탐색해야 한다. pq에 더 이상 간선이 남아있지 않다면 dest 배열에는 각 정점까지 방문하는데 필요한 최단 경로의 값들이 저장된 우선순위 큐가 남게 된다. 만약 모든 간선을 다 탐색했을때 그 정점까지 가는 경로가 K개가 되지 않는다면 그 정점으로 가기위한 K번째 최단 경로가 없다고 판단하고 -1을 출력하면 된다. 경로가 K개라면 peek값을 출력하여 현재 정점 중 가장 큰 값이자 K번째로 작은 최단 경로 값을 출력하면 된다.

 

예시를 분석해보면 1번에서 1번 정점으로 가는 2 번째 최단 경로는 4, 1에서 2번 정점으로 가는 2 번째 최단 경로는 4, 1에서 3번 정점으로 가는 2 번째 최단 경로는 5가 됨을 알 수 있다.

1 → 4

2 → 4

3 → 5

 

이를 구현한 코드는 다음과 같다.

코드

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

public class Main {
    static int N, M, K;
    static ArrayList<Edge>[] graph;

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        StringBuilder sb = new StringBuilder();

        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());

        dest = new PriorityQueue[N + 1];
        graph = new ArrayList[N + 1];

        for (int i = 1; i <= N; i++) {
            dest[i] = new PriorityQueue<>(Comparator.reverseOrder());
            graph[i] = new ArrayList<>();
        }

        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int from = Integer.parseInt(st.nextToken());
            int to = Integer.parseInt(st.nextToken());
            int weight = Integer.parseInt(st.nextToken());
            graph[from].add(new Edge(to, weight));
        }

        findPath(1);

        for (int i = 1; i <= N ; i++) {
            if (dest[i].size() < K) {
                sb.append("-1\\n");
            } else {
                sb.append(dest[i].peek()).append("\\n");
            }
        }
        System.out.print(sb);
    }

    public static void findPath(int root) {
        PriorityQueue<Edge> pq = new PriorityQueue<>();
        dest[root].add(0);
        pq.add(new Edge(root, 0));

        while (!pq.isEmpty()) {
            Edge current = pq.poll();
            for (Edge next: graph[current.to]) {
                if (dest[next.to].size() < K || dest[next.to].peek() > current.weight + next.weight) {
                    if (dest[next.to].size() == K) dest[next.to].poll();
                    dest[next.to].add(current.weight + next.weight);
                    pq.add(new Edge(next.to, current.weight + next.weight));
                }
            }
        }
    }
}

class Edge implements Comparable<Edge> {
    int to;
    int weight;

    public Edge(int to, int weight) {
        this.to = to;
        this.weight = weight;
    }

    @Override
    public int compareTo(Edge e) {
        return this.weight - e.weight;
    }
}