부대복귀(프로그래머스-연습문제)

문제

코드

import java.util.*;
class Solution {
    public int[] solution(int n, int[][] roads, int[] sources, int destination) {
        int[] answer = new int[sources.length];
        ArrayList<Integer>[] graph=new ArrayList[n+1];
        for(int i=1;i<=n;i++)
            graph[i]=new ArrayList();
        
        for(int[] road:roads){
            int start=road[0];
            int end=road[1];
            graph[start].add(end);
            graph[end].add(start);
        }
        int[] distance=dijkstra(n,graph,destination);
        for(int i=0;i<answer.length;i++){
            if(distance[sources[i]]==0){
                if(sources[i]==destination) answer[i]=0;
                else answer[i]=-1;
            }
            else answer[i]=distance[sources[i]];
        }
        return answer;
    }
    public int[] dijkstra(int n, ArrayList<Integer>[] graph, int destination){
        Queue<Node> queue=new LinkedList<>();
        queue.add(new Node(destination,0));
        int[] distanceList=new int[n+1];
        boolean[] visited=new boolean[n+1];
        
        while(!queue.isEmpty()){
            Node start=queue.poll();
            visited[start.position]=true;
            for(int node:graph[start.position]){
                if(!visited[node]){
                    distanceList[node]=start.distance+1;
                    queue.add(new Node(node,start.distance+1));
                    visited[node]=true;
                }
            }
        }
        return distanceList;
    }
}
class Node{
    int position;
    int distance;
    Node(int position, int distance){
        this.position=position;
        this.distance=distance;
    }
}

풀이

이 문제는 인접리스트를 이용한 다익스트라 알고리즘을 통해 해결할 수 있다.
이 문제의 핵심은 각 지역에서 도착지까지의 거리를 일일히 구하는 것이 아닌,
도착지에서 각 지역까지의 거리를 구해주는 것이다.
만약, 각 지역마다 도착지까지의 거리를 구하는 작업을 일일히 한다면 아마 시간초과가 날 것이다.
하지만, 도착지에서 출발해서 각 지역까지 거리를 저장해준다면 roads 배열을 한번만 돌면 된다.

(1) 인접 리스트를 만들어준다. 인접리스트에는 각 지역에서 연결된 모든 지역이 들어간다.
(2) roads 배열을 돌면서, 인접 리스트를 채워준다.
(3) dijkstra() 함수를 실행하여 각 지역마다 최소 거리 정보가 담긴 배열을 만든다.
(4) dijkstra() 함수에서는 destination 에서 출발하여, 각 지역까지의 최단 거리를 구해준다.

• 큐를 생성한다. 큐에는 Node 클래스가 들어갈 것이다. Node 클래스에는 각 지역(position)과 distance(destination 까지 최단 거리)
• 큐에 destination 지역 노드를 추가하고(distance는 0) 반복문을 실행한다.
• 만들어 놓은 인접리스트를 통해 큐에서 빼낸 노드에서 갈 수 있는 지역을 확인한다. 만약 방문하지 않은 노드라면, 큐에 추가하고 최단 거리 리스트를 갱신한다.
• 위 과정을 큐가 비어있을 때까지 반복한다.
  (5) dijkstra() 함수에서 생성한 최단 거리 리스트를 통해, answer에 sources 배열 안의 각 지역마다의 최단 거리를 넣어준다.(최단 거리가 0일 경우에, 그 지역이 destination이 아니면 -1을 넣어준다.)

결과

처음에 나는, 목적지 부터 출발해야한다는 생각을 못했을뿐더러 인접 리스트로 풀어야 한다는 생각조차 하지 못했다.
처음엔 그냥 각 목적지 에서 출발해서 bfs 탐색을 통해 destination을 찾았는데, 대부분의 케이스에서 시간초과가 발생했다.

이런 양방향 그래프 문제는 익숙하지 않아서 생각하기가 조금 어려웠다.
하지만 다음에 이런 문제가 나오면 헤매지 않아야 한다.

출처 : 프로그래머스 코딩 테스트 연습 https://school.programmers.co.kr/learn/challenges