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

사용 언어: Kotlin

 

문제

https://school.programmers.co.kr/learn/courses/30/lessons/132266

프로그래머스

특정 지역으로 이동해 임무를 수행중인 부대원이 부대의 지역으로 복귀할 수 있는지, 각 지역에서의 복귀할때의 최단시간을 return한다.

복귀가 불가능하면 -1을 return한다.

 

제한사항

n = 지역의 개수. 각 지역은 1부터 n까지 번호로 구분한다.

3 <= n <= 100,000

 

roads[a,b] = a에서 b로 왕복할 수 있는 길이 있다.

2 <= roads의 길이 <= 500,000

동일한 정보가 중복해서 주어지지 않는다.

 

sources[i] = 부대원이 i 위치에 있다.

1 <= sources의 길이 <= 500

1 <= sources[i] <= n

 

destination = 부대의 위치

1 <= destination <= n

풀이

기본적인 DFS/BFS 문제이다.

destination을 루트로 하여 BFS 순회하며 각 지역의 거리를 측정한다.

탐색 후, 거리가 측정되지 않은 노드는 부대 복귀가 불가능한 것으로 판단하여 -1을 갖도록 한다.

 

코드

fun solution(n: Int, roads: Array<IntArray>, sources: IntArray, destination: Int): IntArray {
    val edge = Array(n+1) { arrayListOf<Int>() }
    val distance = IntArray(n+1) { -1 }
    val visit = BooleanArray(n+1)
    distance[destination] = 0

    roads.forEach {
        edge[it[0]].add(it[1])
        edge[it[1]].add(it[0])
    }

    val queue = ArrayDeque<Pair<Int, Int>>()
    queue.add(destination to 0)

    while(queue.isNotEmpty()) {
        val (node, cost) = queue.removeFirst()
        if(visit[node]) continue
        visit[node] = true

        distance[node] = cost
        edge[node].forEach {
            queue.add(it to cost + 1)
        }
    }

    return IntArray(sources.size) { distance[sources[it]] }
}