【프로그래머스】 배달 (Kotlin)

문제

프로그래머스

 

풀이

우선순위 큐를 이용한 다익스트라 알고리즘을 적용했다.

 

그래프

HashMap을 사용하여 각 마을 사이의 도로 정보를 나타내는 그래프를 구성한다.

각 마을 번호를 키, 연결된 마을과 소요 시간을 Pair로 짝짓고 이를 포함하는 MutableList를 값으로 한다.

val graph = hashMapOf<Int, MutableList<Pair<Int, Int>>>()
road.forEach {
    // MutableList에 도로로 연결된 두 마을의 연결 정보를 추가한다.
    graph.getOrPut(it[0]) { mutableListOf() }.add(it[1] to it[2])
    graph.getOrPut(it[1]) { mutableListOf() }.add(it[0] to it[2])
}

 

우선순위 큐

마을과 해당 마을까지 걸리는 시간을 짝짓는 Pair 객체를 정렬 상태로 유지하기 위해 PriorityQueue를 사용한다.

// 시간에 대해 오름차순으로 정렬한다.
val comparator = compareBy<Pair<Int, Int>> { it.second }

// 1번 마을은 소요 시간이 없다.
val heap = PriorityQueue(comparator).apply { add(1 to 0) }

 

최단 시간 테이블

IntArray에 각 마을까지 걸리는 최단 시간을 저장한다.

 

1번 마을(인덱스 0)은 소요 시간이 없으므로 0으로 초기화한다.

마을의 개수는 최대 50, 도로를 지나는 데 걸리는 시간은 최대 10000이므로 한 마을까지 걸리는 최대 시간은 500000이다.

따라서 나머지 인덱스의 값은 무한 대신 500001으로 초기화한다.

val table = IntArray(N) { if (it == 0) 0 else 500001 }

 

반복문

heap에서 꺼낸 데이터를 table의 데이터와 비교하여 table을 업데이트한다.

while (heap.isNotEmpty()) {
    val current = heap.poll()
    val town = current.first  // 마을
    val time = current.second // town까지 걸리는 시간
    
    // time이 table에 저장된 시간보다 크다면 갱신할 필요가 없다.
    if (time <= table[town - 1]) {
    
        // town에 연결된 각 도로에 대해
        for (info in graph[town]!!) {
            val neighbor = info.first        // 도로로 연결된 마을
            val newTime = time + info.second // time + 도로를 지나는 데 걸리는 시간
            
            // newTime이 table에 저장된 시간보다 작다면 갱신한다.
            // 최단 경로를 이후에 이어서 탐색하기 위해 heap에 데이터를 추가한다.
            if (newTime < table[neighbor - 1]) {
                table[neighbor - 1] = newTime
                heap.add(neighbor to newTime)
            }
        }
    }
}

 

마지막으로 table에서 k 이하인 요소의 개수를 반환한다.

return table.count { it <= k }

 

코드

GitHub - blacksg/ProblemSolving

import java.util.PriorityQueue

class Solution {
    fun solution(N: Int, road: Array<IntArray>, k: Int): Int {
        val graph = hashMapOf<Int, MutableList<Pair<Int, Int>>>()
        road.forEach {
            graph.getOrPut(it[0]) { mutableListOf() }.add(it[1] to it[2])
            graph.getOrPut(it[1]) { mutableListOf() }.add(it[0] to it[2])
        }
        val comparator = compareBy<Pair<Int, Int>> { it.second }
        val heap = PriorityQueue(comparator).apply { add(1 to 0) }
        val table = IntArray(N) { if (it == 0) 0 else 500001 }
        while (heap.isNotEmpty()) {
            val current = heap.poll()
            val town = current.first
            val time = current.second
            if (time <= table[town - 1]) {
                for (info in graph[town]!!) {
                    val neighbor = info.first
                    val newTime = time + info.second
                    if (newTime < table[neighbor - 1]) {
                        table[neighbor - 1] = newTime
                        heap.add(neighbor to newTime)
                    }
                }
            }
        }
        return table.count { it <= k }
    }
}