[백준] 특정 거리의 도시 찾기 (18352)(kotlin)

문제 설명

백준 18352번 문제 링크

입력 및 출력

» 입력

첫째 줄에 도시의 개수 N, 도로의 개수 M, 거리 정보 K, 출발 도시의 번호 X가 주어진다. (2 ≤ N ≤ 300,000, 1 ≤ M ≤ 1,000,000, 1 ≤ K ≤ 300,000, 1 ≤ X ≤ N) 둘째 줄부터 M개의 줄에 걸쳐서 두 개의 자연수 A, B가 공백을 기준으로 구분되어 주어진다. 이는 A번 도시에서 B번 도시로 이동하는 단방향 도로가 존재한다는 의미다. (1 ≤ A, B ≤ N) 단, A와 B는 서로 다른 자연수이다.

» 출력

• X로부터 출발하여 도달할 수 있는 도시 중에서, 최단 거리가 K인 모든 도시의 번호를 한 줄에 하나씩 오름차순으로 출력한다. 이 때 도달할 수 있는 도시 중에서, 최단 거리가 K인 도시가 하나도 존재하지 않으면 -1을 출력한다.

예제 입출력(테스트케이스)

입력	출력
4 4 2 1 1 2 1 3 2 3 2 4	4
4 3 2 1 1 2 1 3 1 4	-1
4 4 1 1 1 2 1 3 2 3 2 4	2 3

문제 풀이1

import java.util.PriorityQueue

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

fun main() = with(System.`in`.bufferedReader()) {
    val (N, M, K, X) = readLine().split(" ").map { it.toInt() }
    val adjMatrix = Array(N + 1) { PriorityQueue<Node>() }

    repeat(M) {
        val (A, B) = readLine().split(" ").map { it.toInt() }
        adjMatrix[A].add(Node(B, 1))
    }

    val result = dijkstra(X, N, K, adjMatrix)
    printResult(K, result)
}

fun printResult(K: Int, result: IntArray) {
    val list = mutableListOf<Int>()

    for (i in result.indices) {
        if (result[i] == K) {
            list.add(i)
        }
    }

    if (list.isNotEmpty()) {
        list.forEach { println(it) }
    } else {
        println(-1)
    }
}

fun dijkstra(start: Int, N: Int, K: Int, adjMatrix: Array<PriorityQueue<Node>>): IntArray {
    val check = BooleanArray(N + 1)
    val distance = IntArray(N + 1) { Int.MAX_VALUE }
    val q = PriorityQueue<Node>()

    distance[start] = 0
    q.add(Node(start, 0))

    while (q.isNotEmpty()) {
        val now = q.poll()

        if (check[now.destination]) continue
        check[now.destination] = true

        for (next in adjMatrix[now.destination]) {
            if (distance[next.destination] > distance[now.destination] + next.cost) {
                distance[next.destination] = distance[now.destination] + next.cost
                q.add(Node(next.destination, distance[next.destination]))
            }
        }
    }

    return distance
}