[JAVA] 부대복귀

 

부대복귀

링크

프로그래머스

 

난이도 : Level 3

유형 : 최단거리 찾기

 

문제 설명

강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다.
지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 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

풀이

정점의 개수는 최대 100,000개, 목적지는 정해져있고 출발지는 최대 500개인 문제다. 

roads의 조건을 보면 왕복이 가능하다했으니 양방향 그래프이고 가중치는 딱히 주어지지 않았으니 1로 두면 될 것 같다.

그리고 음수 가중치가 없으니 다익스트라를 써서 최단거리를 찾으면 될 것 같다.

 

간단하게 생각하면 그냥 각 출발지에서 도착지까지의 거리를 찾아서 result 배열에 넣으면 끝아닌가? 라고 생각할 수 있다.

그러면 출발지의 개수가 n개 라고 하면 n번의 다익스트라를 해야한다. 이건 너무 많은거같다.

 

살짝 거꾸로 생각해보면 1번의 다익스트라로도 가능하다.

도착지에서 다익스트라로 각 정점까지의 최단거리를 구하면 그게 각 출발지에서 도착지까지 최단거리들의 집합과도 같다.

따라서 1번의 다익스트라로 모든 값을 구할 수 있는 것이다. 

 

아마 그냥 각 출발지마다 다익스트라 돌리면 시간초과가 날 것 같다.

근데 처음엔 그래프를 습관대로 인접행렬로 표현했는데 이러니까 메모리초과가 난다;

인접 행렬의 공간 복잡도는 O(n^2), 인접 리스트의 공간복잡도는 O(n+m)이다. (m은 간선의 수)

정점의 수가 많을땐 인접리스트형태를 쓰자.

 

코드

 

import java.util.*;

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

        // 인접 리스트 초기화
        List<List<Integer>> graph = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            graph.add(new ArrayList<>());
        }

        // 간선 정보를 인접 리스트에 반영
        for (int[] road : roads) {
            int u = road[0] - 1; // 0-based 인덱스로 변환
            int v = road[1] - 1;
            graph.get(u).add(v);
            graph.get(v).add(u); // 양방향 그래프
        }

        // 다익스트라 알고리즘을 사용하여 destination에서 모든 노드까지의 최단 거리 계산
        int[] distances = dijkstra(graph, destination - 1, n);

        // sources에 대한 결과 계산
        for (int i = 0; i < sources.length; i++) {
            int source = sources[i] - 1; // 0-based 인덱스로 변환
            if (distances[source] != Integer.MAX_VALUE) {
                answer[i] = distances[source];
            } else {
                answer[i] = -1;
            }
        }

        return answer;
    }

    // 우선순위 큐를 사용한 다익스트라 알고리즘
    private int[] dijkstra(List<List<Integer>> graph, int start, int n) {
        int[] distances = new int[n]; // 시작 노드에서 각 노드까지의 최단 거리
        Arrays.fill(distances, Integer.MAX_VALUE);
        distances[start] = 0;

        // 우선순위 큐: (거리, 노드)를 저장
        PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));
        pq.offer(new int[]{0, start});

        while (!pq.isEmpty()) {
            int[] current = pq.poll();
            int dist = current[0];
            int u = current[1];

            // 현재 노드의 거리가 이미 더 작은 경우 무시
            if (dist > distances[u]) continue;

            // 이웃 노드 탐색
            for (int v : graph.get(u)) {
                int newDist = dist + 1; // 모든 간선의 가중치는 1
                if (newDist < distances[v]) {
                    distances[v] = newDist;
                    pq.offer(new int[]{newDist, v});
                }
            }
        }

        return distances;
    }
}