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

문제

강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며,

두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다.

다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.

강철부대가 위치한 지역을 포함한 총지역의 수 n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열 roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열 sources,

강철부대의 지역 destination이 주어졌을 때, 주어진 sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요.

복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.

풀이

지역별 탐색을 해서 최단 거리를 구해야 하는 것이다.
BFS를 이용하지, 다익스트라를 이용할지, 플로이드 워셜을 이용할지 고민었다.
다익스트라를 연습할 겸 우선순위를 이용한 다익스트라를 이용하여 풀이를 하였다.
시간상으로는 BFS로 풀면 O(V + e)로 시간이 더 빠를 것이라고 추측된다.

처음 풀이를 다익스트라르 여러번 시도하는 것이었다.
sources원소를 하나씩 다익스트라를 시도하면서 거리를 구하였다.
원소를 돌릴때 마다 distance 배열과 방문 배열을 초기화 시켜야 하기 때문에 시간초과가 나왔다.

생각을 한번 뒤집어보자.
전체 원소를 하나씩 돌리는것이 아니라 destination을 기준으로 한번만 다익스트라를 돌리면 출발지에 대한 모든 정점
비용을 얻을 수 있다. 따라서 더 효율적인 풀이가 가능하였다.

import java.util.*;
import java.lang.*;

class Solution {
    static int[] distance;
    static List<Node>[] list;
    static boolean[] visited;
    static int MAX_INT = 1000000;
    public int[] solution(int n, int[][] roads, int[] sources, int destination) {
        int[] answer = new int[sources.length];
        list = new List[n + 1];
        for(int i = 1; i <= n; i++) {
            list[i] = new ArrayList<>();
        }

        for(int i = 0; i < roads.length; i++) {
            int[] road = roads[i];
            list[road[0]].add(new Node(road[1], 1));
            list[road[1]].add(new Node(road[0], 1));
        }
        distance = new int[n + 1];
        visited = new boolean[n + 1];
        Arrays.fill(distance, MAX_INT);

        dikstra(destination);

        for(int i = 0; i < sources.length; i++) {
            if (distance[sources[i]] == MAX_INT) {
                answer[i] = -1;
                continue;
            }

            answer[i] = distance[sources[i]];
        }


        return answer;
    }

    static void dikstra(int start) {
        Queue<Node> queue = new PriorityQueue<>((o1, o2) -> o1.cost - o2.cost);
        distance[start] = 0;
        queue.add(new Node(start, 0));

        while(!queue.isEmpty()) {
            Node currentNode = queue.poll();

            if (visited[currentNode.des]) {
                continue;
            }

            visited[currentNode.des] = true;

            for (Node node : list[currentNode.des]) {
                if (distance[node.des] > node.cost + distance[currentNode.des]) {
                    distance[node.des] = node.cost + distance[currentNode.des];
                    queue.add(new Node(node.des, distance[node.des]));
                }
            }

        }

    }

    static class Node {
        int des;
        int cost;

        public Node(int des, int cost) {
            this.des = des;
            this.cost = cost;
        }
    }
}