코딩 테스트 [그래프] - K번째 최단 경로 찾기

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

(시간 제한 2초)

---

입력

1번째 줄에 n, m, k가 주어진다.(1 ≤ N ≤ 1,000, 0 ≤ m ≤ 2,000,000, 1 ≤ k ≤ 100) n과 m은 각각 김 조교가 여행을 고려하고 있는 도시들의 개수와 도시간에 존재하는 도로의 수다. 이어지는 m개의 줄에는 각각 도로의 정보를 제공하는 3개의 정수 a, b, c가 포함돼 있다. 이것은 a번 도시에서 b번 도시로 갈 때는 c의 시간이 걸린다는 의미다(1 ≤ a,b ≤ n, 1 ≤ c ≤ 1,000). 도시의 번호는 1번부터 n번까지 연속해 있고, 1번은 시작 도시다.

출력

n개의 줄을 출력한다. i번째 줄에 1번 도시에서 i번 도시로 가는 K번째 최당 경로의 소요 시간을 출력한다. 경로의 소요 시간은 경로 위에 있는 도로들을 따라 이동하는데 필요한 시간들의 합이다. i번 도시에서 i번 도시로 가는 최단 경로는 0이지만, 일반적인 K번째 최단 경로는 0이 아닐 수 있다는 것에 유의한다. 또 K번째 최단 경로가 존재하지 않으면 -1을 출력한다. 최단 경로에 같은 노드가 여러 번 포함될 수 있다.

예제 입력
5 10 2	// 도시 수, 도로 수, K
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

예제 출력
-1
10
7
5
14

---

1단계 - 문제 분석하기

시작점과 도착점이 주어지고 이 목적지까지 가는 K번째 최단 경로를 구하는 문제다. 도시(노드)의 개수는 1,000개. 도로(에지)의 수는 2,000,000이면서 시간 제약이 2초이므로 다익스트라 알고리즘으로 접근할 수 있다.
이 문제에서 가장 고민되는 부분은 최단 경로가 아닌 K번째 최단 경로라는 점이다. 이 부분을 해결하기 위해 다음과 같이 변경해 보려고 한다.

K번째 최단 경로 해결 방법

• 최단 경로를 표현하는 배열을 우선순위 큐 배열(크기는 K)로 변경하고자 한다. 이러면 최단 경로뿐 아니라 최단 경로 ~ K번째 최단 경로까지 표현할 수 있다.
• 사용한 노드는 방문 배열에 확인해 두고 재사용하지 않는다는 부분은 삭제한다. K번째 경로를 찾기 위해서는 노드를 여러 번 쓰는 경우가 생기기 때문이다.

2단계 - 손으로 풀어 보기

1 주어진 데이터를 기반으로 그래프를 그린다.

2 변수를 선언하고 그래프 데이터를 받는 부분을 다익스트라 알고리즘 준비 과정과 동일하게 준비한다.

3 최단 거리 배열을 우선순위 큐 배열로 선언하고, 다음과 같은 기준을 세워 채워야 한다.
예제에서 주어진 입력에서는 K = 2이지만, 이해를 돕기 위해 K = 3인 경우를 계산해보겠다.

최단 거리 배열 채우기 규칙

1. 현재 노드에 저장된 경로가 K개 미만일 때 신규 경로를 추가한다.
2. 경로가 K개일 때 현재 경로 중 최대 경로와 신규 경로를 비교해 신규 경로가 더 작을 때 업데이트한다. 우선순위 큐를 사용하면 이 로직을 쉽게 구현할 수 있다.
3. K번째 경로를 찾기 위해서는 노드를 여러 번 쓰는 경우가 생기므로 사용한 노드는 방문 배열에 확인해 두고 재사용하지 않는 부분은 삭제한다.

4 최단 거리 배열을 탐색하면서 K번째 경로가 존재하면 출력하고, 존재하지 않으면 -1을 출력한다.

3단계 - sudo코드 작성하기

N(노드 개수) M(에지 개수) K(목표 최단 경로)
W(그래프 저장 인접 행렬)
distQueue(최단 거리 우선순위 큐)

for(노드 개수 + 1)
{
	최단 거리 큐 배열 초기화
}

for(에지 개수)
{
	인접 행렬에 에지 정보 저장
}

pq(우선순위 큐 선언)
출발 노드를 우선순위 큐에 삽입
while(큐가 빌 때까지)
{
	현재 노드 = 우선순위 큐에서 poll
	for(노드 개수만큼 반복)
    {
    	if(해당 노드가 현재 노드와 연결돼 있으면)
        {
        	if(최단 거리 배열 큐에 해당 노드에 관해 저장된 경로가 K보다 작으면)
            {
            	최단 거리 큐 배열에 거리 정보 삽입하고 큐에 선택 노드를 추가하기
            }
            else if(최단 거리 큐의 마지막 값 > 이전 노드의 값 + 두 노드 사이의 에지 가중치)
            {
            	해당 노드에 최단 거리 큐에 마지막값 삭제하고 신규 값으로 업데이트
                큐에 선택 노드 추가하기
            }
        }
    }
}

for(노드 개수)
{
	우선순위 큐 크기가 K이면 큐의 값 출력, 아니면 -1 출력
}

4단계 - 코드 구현하기

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class Q58 {
    public static void main(String[] args) throws IOException {
        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());

        int[][] W = new int[1001][1001];
        // ArrayList<Node58>[] W = new ArrayList[N + 1];

        PriorityQueue<Integer>[] distQueue = new PriorityQueue[N+1];
        Comparator<Integer> cp = new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o1 < o2 ? 1 : -1;
            }
        };

        for(int i = 1; i < N+1; i++){
            distQueue[i] = new PriorityQueue<Integer>(K, cp);
            // W[i] = new ArrayList<Node58>;
        }

        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 V = Integer.parseInt(st.nextToken());
            W[S][E] = V;
            // W[S].add(new Node58(E, V));
        }

        PriorityQueue<Node58> pq = new PriorityQueue<>();
        pq.add(new Node58(1, 0));
        distQueue[1].add(0);
        while (!pq.isEmpty()){
            Node58 nowNode = pq.poll();
            for(int nextNode = 1; nextNode < N+1; nextNode++){
                if(W[nowNode.node][nextNode] != 0){             // for(int i = 0; i < W[nowNode].size(); i++) {
                    if(distQueue[nextNode].size() < K){
                        distQueue[nextNode].add(nowNode.value + W[nowNode.node][nextNode]);
                        pq.add(new Node58(nextNode, nowNode.value + W[nowNode.node][nextNode]));
                    } else if (distQueue[nextNode].peek() > nowNode.value + W[nowNode.node][nextNode]) {
                        distQueue[nextNode].poll();
                        distQueue[nextNode].add(nowNode.value + W[nowNode.node][nextNode]);
                        pq.add(new Node58(nextNode, nowNode.value + W[nowNode.node][nextNode]));
                    }
                }
            }
        }

        for(int i = 1; i < N+1; i++){
            if(distQueue[i].size() == K){
                System.out.println(distQueue[i].poll());
            }else {
                System.out.println(-1);
            }
        }

        br.close();
    }
}

class Node58 implements Comparable<Node58> {
    int node;
    int value;
    Node58(int node, int value){
        this.node = node;
        this.value = value;
    }

    @Override
    public int compareTo(Node58 o) {
        return this.value < o.value ? -1 : 1;
    }
}

---

• Do it! 알고리즘 코딩테스트 자바 편