다익스트라 알고리즘 : 파티 [백준 1238, Kotlin]

[문제]

N개의 숫자로 구분된 각각의 마을에 한 명의 학생이 살고 있다.
어느 날 이 N명의 학생이 X (1 ≤ X ≤ N)번 마을에 모여서 파티를 벌이기로 했다. 이 마을 사이에는 총 M개의 단방향 도로들이 있고 i번째 길을 지나는데 Ti(1 ≤ Ti ≤ 100)의 시간을 소비한다.
각각의 학생들은 파티에 참석하기 위해 걸어가서 다시 그들의 마을로 돌아와야 한다. 하지만 이 학생들은 워낙 게을러서 최단 시간에 오고 가기를 원한다.
이 도로들은 단방향이기 때문에 아마 그들이 오고 가는 길이 다를지도 모른다. N명의 학생들 중 오고 가는데 가장 많은 시간을 소비하는 학생은 누구일지 구하여라.

 

[입력 조건]

• 첫째 줄에 N(1 ≤ N ≤ 1,000), M(1 ≤ M ≤ 10,000), X가 공백으로 구분되어 입력된다. 두 번째 줄부터 M+1번째 줄까지 i번째 도로의 시작점, 끝점, 그리고 이 도로를 지나는데 필요한 소요시간 Ti가 들어온다. 시작점과 끝점이 같은 도로는 없으며, 시작점과 한 도시 A에서 다른 도시 B로 가는 도로의 개수는 최대 1개이다.
• 모든 학생들은 집에서 X에 갈수 있고, X에서 집으로 돌아올 수 있는 데이터만 입력으로 주어진다

 

[출력 조건]

• 첫 번째 줄에 N명의 학생들 중 오고 가는데 가장 오래 걸리는 학생의 소요시간을 출력한다.

 

처음으로 해결한 방법

import java.io.BufferedReader
import java.io.InputStreamReader
import java.util.*

private lateinit var result : IntArray

// 다른 마을에서 x 마을까지 최단 거리 + x 마을에서 다른 마을까지 최단 거리
fun main(){
    val br = BufferedReader(InputStreamReader(System.`in`))
    val (n, m, x) = br.readLine().split(" ").map{ it.toInt() }
    val graph = Array(n+1){ mutableListOf<Pair<Int, Int>>() }
    result = IntArray(n + 1)

    // 그래프 입력
    repeat(m){
        val (a, b, c) = br.readLine().split(" ").map{ it.toInt() }
        graph[a].add(Pair(b, c))
    }

    val distance = (0..n).map{ Int.MAX_VALUE }.toMutableList()

    // 간선 값으로 정렬된 우선순위 큐
    val q = PriorityQueue<Pair<Int, Int>>{ p1, p2 ->
        when {
            p1.second > p2.second -> 1
            p1.second < p2.second -> -1
            else -> 0
        }
    }

    q.add(Pair(x, 0)) // 시작 마을, 간선 값 0
    distance[x] = 0

    while(q.isNotEmpty()){
        val (now, dist) = q.poll() // 최단거리가 가장 짧은 값 꺼내기

        if (distance[now] < dist) // 처리된 값이면 무시
            continue

        for(i in graph[now]){
            val cost = dist + i.second
            if(cost < distance[i.first]){ // 거치는 값이 더 작은 경우
                distance[i.first] = cost
                q.add(Pair(i.first, cost))
            }
        }
    }

    distance.forEachIndexed { index, value ->
        if(value != Int.MAX_VALUE){
            result[index] += value
        }
    }

    var max = - 1
    for(i in 1 .. n){
        val temp = dijkstra(i, graph, n, x)
        if(temp != Int.MAX_VALUE){
            result[i] += temp
            max = maxOf(max, result[i])
        }
    }
    print(max)
}

// 모든 마을에서 x 마을까지 거리
fun dijkstra(start: Int, graph: Array<MutableList<Pair<Int, Int>>>, n : Int, x: Int): Int{
    val distance = (0..n).map{ Int.MAX_VALUE }.toMutableList()

    val q = PriorityQueue<Pair<Int, Int>>{ p1, p2 ->
        when {
            p1.second > p2.second -> 1
            p1.second < p2.second -> -1
            else -> 0
        }
    }

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

    while(q.isNotEmpty()){
        val (now, dist) = q.poll()
        if(now == x){
            return dist
        }
        if(distance[now] < dist)
            continue
        for(i in graph[now]){
            val cost = dist + i.second
            if(cost < distance[i.first]){
                distance[i.first] = cost
                q.add(Pair(i.first, cost))
            }
        }
    }
    return Int.MAX_VALUE // 갈수 없는 경우
}

처음으로 해결한 방법은 왕복 거리를 구하기 위해서 x 마을에서 다른 마을까지의 거리와 다른 마을에서 x 마을까지의 거리의 합을 다익스트라 알고리즘을 통해 구현하였다. 해당 풀이는 모든 마을에서 x 마을까지 다익스트라를 탐색하므로 지금의 조건과 다르게 N이 커진다면 시간초과가 발생할 수 있는 문제이다.

 

두번째로 해결한 방법 

다익스트라 알고리즘은 하나의 시작점으로 부터 다른 모든 정점의 거리를 구할 수 있는 알고리즘이다. 방향 그래프에서 이를 반대로 생각하여, 간선의 방향이 반대로 바꾼다면 (a -> b를 b -> a로) 다른 모든 정점으로 부터 하나의 시작점까지의 거리를 구할 수 있다. 

import java.io.BufferedReader
import java.io.InputStreamReader
import java.util.*

private lateinit var graph : Array<MutableList<Pair<Int, Int>>>
private lateinit var reverseGraph : Array<MutableList<Pair<Int, Int>>>
private var n = 0
private var x = 0

fun main(){
    val br = BufferedReader(InputStreamReader(System.`in`))
    val stk = StringTokenizer(br.readLine())
    n = stk.nextToken().toInt()
    val m = stk.nextToken().toInt()
    x = stk.nextToken().toInt()
    graph = Array(n + 1){ mutableListOf() }
    reverseGraph = Array(n + 1){ mutableListOf() }

    repeat(m){
        val (a, b, c) = br.readLine().split(" ").map{ it.toInt() }
        graph[a].add(Pair(b, c)) // 정방향 그래프
        reverseGraph[b].add(Pair(a, c)) // 역방향 그래프
    }
    val distance1 = dijkstra(graph) // x 마을에서 다른 노드까지의 거리
    val distance2 = dijkstra(reverseGraph) // 다른 노드에서 x 마을까지 거리

    var result = 0
    for(i in 1 .. n){
        result = maxOf(result, distance1[i] + distance2[i])
    }
    print(result)
}

fun dijkstra(graph: Array<MutableList<Pair<Int, Int>>>): IntArray{
    val distance = IntArray(n + 1){ Int.MAX_VALUE }
    val visited = BooleanArray(n + 1)
    val q = PriorityQueue<Pair<Int, Int>>{ p1, p2 ->
        when {
            p1.second > p2.second -> 1
            p1.second < p2.second -> -1
            else -> 0
        }
    }
    q.add(Pair(x, 0))
    distance[x] = 0

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

        if(visited[now])
            continue
        visited[now] = true

        for(i in graph[now]){
            val cost = dist + i.second
            if(cost < distance[i.first]){
                distance[i.first] = cost
                q.add(Pair(i.first, cost))
            }
        }
    }
    return distance
}

입력을 받을 때 reverseGraph를 하나 더 만들어 역방향 그래프를 만든다. 정방향 그래프와 역방향를 모두 다익스트라 알고리즘을 적용하여 왕복 최단 거리를 구할 수 있었다. 

 

결론적으로 다익스트라를 매번 노드마다 호출하지 않고, 한번만 호출하게 하여 코드의 길이와 실행 시간이 대폭 감소되었다.