1504번: 특정한 최단 경로 - Kotlin

1504번: 특정한 최단 경로

문제

방향성이 없는 그래프가 주어진다. 세준이는 1번 정점에서 N번 정점으로 최단 거리로 이동하려고 한다. 또한 세준이는 두 가지 조건을 만족하면서 이동하는 특정한 최단 경로를 구하고 싶은데, 그것은 바로 임의로 주어진 두 정점은 반드시 통과해야 한다는 것이다.

세준이는 한번 이동했던 정점은 물론, 한번 이동했던 간선도 다시 이동할 수 있다. 하지만 반드시 최단 경로로 이동해야 한다는 사실에 주의하라. 1번 정점에서 N번 정점으로 이동할 때, 주어진 두 정점을 반드시 거치면서 최단 경로로 이동하는 프로그램을 작성하시오.

입력

첫째 줄에 정점의 개수 N과 간선의 개수 E가 주어진다. (2 ≤ N ≤ 800, 0 ≤ E ≤ 200,000) 둘째 줄부터 E개의 줄에 걸쳐서 세 개의 정수 a, b, c가 주어지는데, a번 정점에서 b번 정점까지 양방향 길이 존재하며, 그 거리가 c라는 뜻이다. (1 ≤ c ≤ 1,000) 다음 줄에는 반드시 거쳐야 하는 두 개의 서로 다른 정점 번호 v1과 v2가 주어진다. (v1 ≠ v2, v1 ≠ N, v2 ≠ 1) 임의의 두 정점 u와 v사이에는 간선이 최대 1개 존재한다.

출력

첫째 줄에 두 개의 정점을 지나는 최단 경로의 길이를 출력한다. 그러한 경로가 없을 때에는 -1을 출력한다.

 

 

나의 풀이

특정한 한 노드에서 다른 모든 노드까지의 최단거리를 구하는 문제임으로 이 문제를 보면 바로 다익스트라가 떠올라야 한다. 기존의 다익스트라 문제랑 거의 다를건 없다. 다만 기존에는 1번 노드부터 N번 노드까지의 최단거리를 구했다면 이번에는 2가지로 경로를 나눠서 문제를 해결해야 한다. 문제에서 주어진대로 우리는 v1과 v2를 무조건 방문해야 한다. 

 

1번 노드에서 v1과 v2를 무조건 방문하고 n을 도착하는 경우의 수는 2가지 밖에 없다

 

경로 1: 시작 -> v1 -> v2 -> n

경로 2: 시작 -> v2 -> v3 -> n

 

이젠 기존에 1번에서 n까지의 최단거리를 계산했던 코드를 함수화 시켜 시작점과 도착점을 주고 각 경로를 다음과 같이 구하면 된다.

 

경로 1: (시작 -> v1의 최단거리) + (v1 -> v2의 최단거리) + (v2 -> n의 최단거리)

경로 2: (시작 -> v2의 최단거리) + (v2 -> v1의 최단거리) + (v1 -> n의 최단거리)

 

추가로 문제에 도착하지 못할 경우 -1을 출력하라 되어있기 때문에 exitProcess(0) 를 사용해 위의 경우에서 하나라도 도착하지 못할 경우 -1을 출력하고 프로그램을 종료하게 해 주었다.

 

import java.util.*
import kotlin.math.min
import kotlin.system.exitProcess

data class Node(val index:Int, val distance:Int):Comparable<Node>{
    override fun compareTo(other:Node):Int =
        distance - other.distance
}

lateinit var  graph: MutableList<PriorityQueue<Node>>

fun main() = with(System.`in`.bufferedReader()) {
    val (n, e) = readLine().split(" ").map { it.toInt() }

    graph = MutableList(n + 1) { PriorityQueue<Node>() }

    repeat(e) {
        val (u, v, w) = readLine().split(" ").map { it.toInt() }
        graph[u].add(Node(v, w))
        graph[v].add(Node(u, w))
    }

    val (v1, v2) = readLine().split(" ").map { it.toInt() }

    val way1 = dijkstra(1, v1, n) + dijkstra(v1, v2, n) + dijkstra(v2, n, n)
    val way2 = dijkstra(1, v2, n) + dijkstra(v2, v1, n) + dijkstra(v1, n, n)

    println(min(way1, way2))
}

fun dijkstra(start: Int, end: Int, n: Int): Int {
    val pq = PriorityQueue<Node>().apply { offer(Node(start, 0)) }
    val result = IntArray(n + 1) { Int.MAX_VALUE }.apply { this[start] = 0 }

    while (pq.isNotEmpty()) {
        val (now, distance) = pq.poll()

        if(distance > result[now]) continue

        for(i in graph[now]){
            val cost = i.distance + result[now]
            val idx = i.index
            if(cost < result[idx]){
                result[idx] = cost
                pq.offer(Node(idx, cost))
            }
        }
    }

    if (result[end] == Int.MAX_VALUE) {
        println(-1)
        exitProcess(0)
    }

    return result[end]
}