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

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

난이도

Level 3

문제

강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다. 다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.

강철부대가 위치한 지역을 포함한 총지역의 수 n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열 roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열 sources, 강철부대의 지역 destination이 주어졌을 때, 주어진 sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요. 복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.

코드

import java.util.*;

class Solution {
    public int[] solution(int n, int[][] roads, int[] sources, int destination) {
        ArrayList<Integer>[] edge = new ArrayList[n+1];
        int[] dis = new int[n + 1];
        for(int i = 0; i <= n; i++){
            dis[i] = Integer.MAX_VALUE;
            edge[i] = new ArrayList<>();
        }
        
        for(int[] r : roads){
            edge[r[0]].add(r[1]);
            edge[r[1]].add(r[0]);
        }
        
        Queue<int[]> q = new LinkedList<>();
        boolean[] visited = new boolean[n + 1];
        q.add(new int[]{destination, 0});
        visited[destination] = true;
        
        while(!q.isEmpty()){
            int[] now = q.poll();
            dis[now[0]] = Math.min(dis[now[0]], now[1]);
            
            for(int next : edge[now[0]]){
                if(visited[next]) continue;
                q.add(new int[]{next, now[1] + 1});
                visited[next] = true;
            }
        }
        
        int[] ans = new int[sources.length];
        for(int i = 0; i < sources.length; i++){
            if(dis[sources[i]] == Integer.MAX_VALUE){
                ans[i] = -1;
            }else{
                ans[i] = dis[sources[i]];
            }
        }
        
        return ans;
    }
}