[백준 1854번] [Kotlin] K번째 최단경로 찾기

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

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

 

난이도 : 플래티넘 4
태그 : 자료 구조, 그래프 이론, 데이크스트라, 우선순위 큐

 

 

설명

다익스트라가 최단 경로를 찾는 알고리즘이라면,

이번엔 이 다익스트라를 이용하여 K 번째 최단 경로를 찾는 문제입니다.

 

해당 문제 같은 경우는 우선순위 큐를 활용할 수 있습니다.

우선순위 큐를 내림차순 정렬로 지정하고,

각 우선순위 큐의 크기가 K개를 넘는다면 하나씩 빼면서 K개를 유지합니다.

 

 

 

소스코드

import java.util.*

data class Node(val to: Int, val weight: Int) : Comparable<Node> {
    override fun compareTo(other: Node): Int {
        return this.weight - other.weight
    }
}

fun main() = with(System.`in`.bufferedReader()) {
    val (n, m, k) = readLine().split(" ").map { it.toInt() }
    val connection = Array(n + 1) { ArrayList<Node>() }

    repeat(m) {
        val st = StringTokenizer(readLine())
        connection[st.nextToken().toInt()].add(Node(st.nextToken().toInt(), st.nextToken().toInt()))
    }

    val answerPQ = dijkstra(n, k, connection, 1)
    val sb = StringBuilder()

    for (i in 1..n) {
        if (answerPQ[i].size == k) {
            sb.append(answerPQ[i].poll())
        } else {
            sb.append("-1")
        }

        sb.append("\n")
    }

    print(sb)
}

fun dijkstra(n: Int, k: Int, connection: Array<ArrayList<Node>>, startNode: Int): Array<PriorityQueue<Int>> {
    val pq = PriorityQueue<Node>()
    pq.offer(Node(startNode, 0))

    val answerPQ = Array(n + 1) { PriorityQueue<Int>(Collections.reverseOrder()) }
    answerPQ[1].add(0)

    while (pq.isNotEmpty()) {
        val target = pq.poll()

        connection[target.to].forEach { nextNode ->
            val newWeight = nextNode.weight + target.weight

            if (answerPQ[nextNode.to].size < k) {
                answerPQ[nextNode.to].add(newWeight)
                pq.add(Node(nextNode.to, newWeight))
            } else if (answerPQ[nextNode.to].peek() > newWeight) {
                answerPQ[nextNode.to].poll()
                answerPQ[nextNode.to].add(newWeight)
                pq.offer(Node(nextNode.to, newWeight))
            }
        }
    }

    return answerPQ
}

 

구현은 일반적인 다익스트라와 크게 다르지 않지만,

다음 노드를 탐색하는 과정에서 우선순위 큐에 있는 원소가 K개 미만이라면 그대로 가중치를 넣고,

만약 K개 이상이고, 가장 큰 값보다 현재의 가중치가 크다면

(우선순위 큐는 내림차순으로 지정되어 있으므로 마지막 원소가 현재 값 보다 크다면)

가장 큰 값을 빼내고 새 값을 추가합니다.

 

 

 

후기

꽤나 머리를 싸맸던 문제였습니다.

다익스트라를 연습하다가, 조금씩 상위 문제를 풀 고 있는데,

역시 플래티넘은 상당히 버겁네요.