[프로그래머스] 부대복귀 with Kotlin

문제

• n: 총 지역 수
• roads: 연결되어 있는 두 지역의 정보
• sources: 각 부대원이 위치한 서로 다른 지역의 정보
• destination: 강철부대가 위치한 지역

두 지역간의 길을 통과하는데 걸리는 시간은 모두 1이다.

sources의 지역에 있는 각 부대원들은 강철 부대로 최단 시간 복귀하고자 한다.

각 부대원들이 복귀까지 걸리는 최단 시간을 반환하는 문제

복귀가 불가능한 부대원은 -1을 반환한다.

 

풀이 과정

1. 각 지역별로 연결되어 있는 타지역의 정보를 저장 (hashMap)
2. destination을 1번 노드로 설정한 후, 각 지역 별 노드를 배열에 저장
   1. 각각의 인덱스가 지역번호가 된다.
   2. 저장은 1번의 hashMap을 이용한 BFS를 사용
3. 노드는 복귀까지의 소요 시간이 되므로, 노드 - 1 == 복귀까지 걸리는 시간이된다.

 

코드

import java.util.*
class Solution {
    // 지역별로 연결되어 있는 타지역
    val hashMap = hashMapOf<Int, ArrayList<Int>>()
    fun solution(n: Int, roads: Array<IntArray>, sources: IntArray, destination: Int): IntArray {
        var answer: IntArray = IntArray(sources.size){-1}
        
        // destination을 1번 노드로 각 지역별 노드 저장소
        val distance = Array(n+1){ if(it == destination) 1 else -1 }
        
        // 지역별 연결되어 있는 타지역 init
        for(road in roads){
            hashMap.getOrPut(road[0]){ arrayListOf<Int>() }.add(road[1])
            hashMap.getOrPut(road[1]){ arrayListOf<Int>() }.add(road[0])
        }
        
        // 지역별 노드 init (BFS)
        val queue: Queue<Pair<Int, Int>> = LinkedList()
        queue.add(destination to 1)
        
        while(queue.isNotEmpty()){
            val now = queue.poll()
            val d = now.first
            val node = now.second
            
            if(hashMap[d] == null) continue
            
            for(i in hashMap[d]!!){
                if(distance[i] != -1) continue
                
                queue.add(i to node+1)
                distance[i] = node+1
            }
        }
        
        // distance의 값은 지역별(인덱스) 노드 값 
        // sorces 별로 1번 노드와의 거리가 answer
        // -1은 연결 X
        sources.forEachIndexed{ idx, source ->
            if(distance[source] != -1){
                answer[idx] = distance[source] - 1
            } 
        }
        
        return answer
    }
}