[백준/BOJ] 1854번: K번째 최단경로 찾기

문제

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

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

해설

다익스트라 알고리즘을 이용해서 풀 수 있는 문제이다.

이 문제에서는 각 노드마다 우선순위 큐를 두어 이 노드를 방문할 때마다의 시간을 저장할 것이다.

val time = Array(cityCnt + 1) { PriorityQueue<Int>(reverseOrder()) }

이 때, 큰 수가 top에 오도록 해야 한다. 즉 내림차순이어야 한다.

새로 구한 정점까지의 소요 시간이 정점의 우선순위 큐의 top보다 작다면, 다익스트라 알고리즘이 종료됬을 때, top이 k번째 최단 시간이 될 수 없다.(코드의 주석에서 추가 설명)

 

해당 노드를 방문할 때 두 종류로 나뉜다.

1. 해당 정점의 우선순위 큐의 크기가 K 이하인 경우(방문을 k번보다 덜 했을 경우)
   여태까지 걸린 시간을 특정 정점의 우선순위 큐에 저장해주고, 다익스트라를 진행하는 전체적인 우선순위 큐에도 저장해준다.
2. 해당 정점의 우선순위 큐의 크기가 K인 경우
   새로 구한 정점까지의 소요 시간이 정점의 우선순위 큐의 top보다 작다면, 정점의 우선순위 큐를 poll()을 한 후, 1번과 같은 과정을 진행한다.

마지막으로 출력을 할 때, 특정 정점의 우선순위 큐의 크기가 K가 아니라면, K번째 최단 경로의 소요시간이 존재하지 않는 것이므로 -1을 출력한다.

소스 코드

import java.util.*
import kotlin.collections.ArrayList

private var cityCnt = 0
private var roadCnt = 0
private var k = 0

data class Travel(val city: Int, val time: Int): Comparable<Travel> {
    override fun compareTo(other: Travel): Int = time - other.time
}

fun main() = with(System.`in`.bufferedReader()) {
    val st = StringTokenizer(readLine())
    cityCnt = st.nextToken().toInt()
    roadCnt = st.nextToken().toInt()
    k = st.nextToken().toInt()

    val graph = Array(cityCnt + 1) { ArrayList<Pair<Int, Int>>() }
    repeat(roadCnt) {
        val (from, to, time) = readLine().split(" ").map { it.toInt() }
        graph[from].add(Pair(to, time))
    }

    val time = dijkstra(graph)
    for(i in 1 .. cityCnt) {
        if(time[i].size != k) println(-1)
        else println(time[i].peek())
    }
}

private fun dijkstra(graph: Array<ArrayList<Pair<Int, Int>>>): Array<PriorityQueue<Int>> {
    val time = Array(cityCnt + 1) { PriorityQueue<Int>(reverseOrder()) }
    val pq = PriorityQueue<Travel>()
    pq.add(Travel(1, 0))
    time[1].add(0)

    while(pq.isNotEmpty()) {
        val cur = pq.poll()
        val curCity = cur.city
        val curTime = cur.time

        for(i in graph[curCity].indices) {
            val nextCity = graph[curCity][i].first
            val nextTime = curTime + graph[curCity][i].second

            if(time[nextCity].size < k) {
                time[nextCity].add(nextTime)
                pq.add(Travel(nextCity, nextTime))
            }
            else if(time[nextCity].peek() > nextTime) {
            // 결론적으로 다익스트라가 종료됬을 때, 가장 적은 시간 두개가 남게 된다. 또한 내림차순이므로 top이 k번째 최단 시간이 된다.
                time[nextCity].poll()
                time[nextCity].add(nextTime)
                pq.add(Travel(nextCity, nextTime))
            }
        }
    }

    return time
}