프로그래머스 3단계 : 부대복귀 (Java 자바)

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

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

 

제한사항

• 3 ≤ n ≤ 100,000
  
  • 각 지역은 정수 1부터 n까지의 번호로 구분됩니다.
• 2 ≤ roads의 길이 ≤ 500,000
  
  • roads의 원소의 길이 = 2
  • roads의 원소는 [a, b] 형태로 두 지역 a, b가 서로 왕복할 수 있음을 의미합니다.(1 ≤ a, b ≤ n, a ≠ b)
  • 동일한 정보가 중복해서 주어지지 않습니다.
    
    • 동일한 [a, b]가 중복해서 주어지지 않습니다.
    • [a, b]가 있다면 [b, a]는 주어지지 않습니다.
• 1 ≤ sources의 길이 ≤ 500
  
  • 1 ≤ sources[i] ≤ n
• 1 ≤ destination ≤ n

 

입출력 예

 

풀이

강철부대의 부대원들이 여러 지역에서 특수 임무를 수행하고, 각 부대원이 최단 시간에 부대로 복귀할 수 있는지를 계산하는 문제입니다. 지도에서의 지역은 번호로 구분되며, 지역 간 이동에는 1의 시간이 소요됩니다. 다만, 복귀 경로가 존재하지 않을 수도 있습니다. 이때, 부대원들의 복귀 가능 여부와 최단 시간을 반환해야 합니다.

 

그래프 생성: 주어진 지역 간 길 정보를 바탕으로 양방향 그래프를 생성합니다. 이 그래프는 map에 저장되며, 각 지역과 이에 연결된 지역들의 목록을 포함합니다.
=> Map<integer, list> map = new HashMap<>();: 그래프를 표현하기 위한 HashMap을 초기화합니다. 이 맵은 지역 번호를 키로 가지며 해당 지역과 연결된 지역 목록을 값으로 갖습니다.

 

< 다익스트라 알고리즘 시작 >
다익스트라 알고리즘: 복귀 가능한 경로를 찾기 위해 다익스트라 알고리즘을 사용합니다. 시작 지점은 destination으로 고정되며, 모든 지역까지의 최단 시간을 계산합니다.

 

boolean[] visited = new boolean[n + 1];: 방문 여부를 나타내는 배열을 초기화합니다. 인덱스 0은 사용하지 않으므로 n + 1로 초기화합니다.

 

int[] dp = new int[n + 1];: 최단 시간을 저장할 배열을 초기화합니다. 초기 값으로는 최댓값 Integer.MAX_VALUE로 설정합니다.

큐 que를 사용하여 BFS를 구현합니다. destination을 큐에 추가하고, 해당 지역부터 출발하여 모든 지역까지의 최단 경로를 계산합니다. 각 지역에 대해 연결된 지역을 확인하고, 아직 방문하지 않은 지역이라면 방문 체크를 하고 최단 경로를 업데이트합니다.

복귀 가능 여부 및 최단 시간 계산: 주어진 sources 배열의 각 부대원에 대해, 해당 지역으로의 최단 시간을 계산하고, 복귀 경로가 없는 경우 -1로 설정합니다.

결과 배열 생성: 모든 부대원에 대한 최단 시간을 계산하고, 이를 result 리스트에 저장합니다.

결과 반환: result 리스트를 배열로 변환하여 최종적인 결과를 반환합니다.

 

for (int[] road : roads) { ... }: roads 배열을 순회하면서 양방향 길을 설정하여 그래프를 만듭니다.

List<Integer> result = new ArrayList<>();: 각 부대원의 복귀 가능 여부와 최단 시간을 저장할 리스트를 초기화합니다.

for (int s : sources) { ... }: sources 배열에 있는 각 부대원에 대해 최단 시간을 계산하고, 복귀 경로가 없는 경우 -1로 설정합니다.

 

import java.util.*;

class Solution {
    public int[] solution(int n, int[][] roads, int[] sources, int destination) {
        int[] answer = {};

        // 그래프를 표현하기 위한 Map 초기화
        Map<Integer, List<Integer>> map = new HashMap<>();

        // roads 배열을 이용하여 그래프 생성
        for (int[] road : roads) {
            int a = road[0];
            int b = road[1];

            // 양방향의 길을 설정합니다.
            List<Integer> a_list = map.getOrDefault(a, new ArrayList<>());
            a_list.add(b);
            map.put(a, a_list);

            List<Integer> b_list = map.getOrDefault(b, new ArrayList<>());
            b_list.add(a);
            map.put(b, b_list);
        }

        // 다익스트라 알고리즘을 사용하여 최단 경로 계산
        boolean[] visited = new boolean[n + 1];
        int[] dp = new int[n + 1];
        
        // dp 배열을 초기화합니다. 최대값으로 설정하여 나중에 비교할 수 있도록 합니다.
        Arrays.fill(dp, Integer.MAX_VALUE);

        // 시작 지점(destination)을 큐에 추가합니다.
        Queue<Integer> que = new LinkedList<>();
        visited[destination] = true;
        que.add(destination);

        dp[destination] = 0;

        while (!que.isEmpty()) {
            int p = que.poll();

            // 현재 지역과 연결된 지역 목록을 가져옵니다.
            List<Integer> p_list = map.get(p);

            // 현재 지역과 연결된 모든 지역들을 순회합니다.
            for (int temp : p_list) {
                // 아직 방문하지 않은 지역이라면
                if (!visited[temp]) {
                    // 방문 체크를 하고, 최단 경로를 업데이트합니다.
                    visited[temp] = true;
                    dp[temp] = dp[p] + 1; // 현재 지역에서 다음 지역으로 이동하는 데 걸리는 시간은 1이므로 1을 더합니다.
                    que.add(temp);
                }
            }
        }

        // 각 부대원의 복귀 가능 여부와 최단 시간을 계산합니다.
        List<Integer> result = new ArrayList<>();
        for (int s : sources) {
            if (dp[s] == Integer.MAX_VALUE) {
                // 복귀 경로가 없는 경우 -1로 설정합니다.
                dp[s] = -1;
            }
            result.add(dp[s]);
        }

        // 결과를 배열로 변환합니다.
        answer = result.stream().mapToInt(i -> i).toArray();

        return answer;
    }
}