[프로그래머스 Level.3] 부대복귀 (연습문제) (Java)

문제 링크

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

프로그래머스

 

 

 

 

코딩테스트 연습 > 연습문제 > 부대복귀

 

 

 

문제 설명

 

 

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

 

---

 

입출력 예

n	roads	sources	destination
3	[[1, 2], [2, 3]]	[2, 3]	1	[1, 2]
5	[[1, 2], [1, 4], [2, 4], [2, 5], [4, 5]]	[1, 3, 5]	5	[2, -1, 0]

 

---

입출력 예 설명

 

 

입출력 예 #1

 

• 지역 2는 지역 1과 길로 연결되어 있기 때문에, 지역 2에서 지역 1의 최단거리는 1입니다.
• 지역 3에서 지역 1로 이동할 수 있는 최단경로는 지역 3 → 지역 2 → 지역 1 순으로 이동하는 것이기 때문에, 지역 3에서 지역 1의 최단거리는 2입니다.
• 따라서 [1, 2]를 return합니다.

 

 

입출력 예 #2

 

• 지역 1에서 지역 5의 최단경로는 지역 1 → 지역 2 → 지역 5 또는 지역 1 → 지역 4 → 지역 5 순으로 이동하는 것이기 때문에, 최단거리는 2입니다.
• 지역 3에서 지역 5로 가는 경로가 없기 때문에, 지역 3에서 지역 5로 가는 최단거리는 -1입니다.
• 지역 5에서 지역 5는 이동할 필요가 없기 때문에, 최단거리는 0입니다.
• 따라서 [2, -1, 0]을 return합니다.

 

 

 

 

 

 

나의 코드

import java.util.*;

// 어차피 목적지는 하나이므로 목적지로부터 모든 지점의 최단거리를 bfs로 구하면 일차원 배열로 표현이 가능
class Solution {
    public int[] solution(int n, int[][] roads, int[] sources, int destination) {
        int[] answer = new int[sources.length];
        int[] dist = new int[n+1];
        boolean[] visited = new boolean[n+1];
        ArrayList<Integer>[] road = new ArrayList[n+1];
        
        Arrays.fill(dist, -1);
        
        for(int i=0; i<=n; i++) 
            road[i] = new ArrayList<>();
        
        for(int i=0; i<roads.length; i++) {
            road[roads[i][0]].add(roads[i][1]);
            road[roads[i][1]].add(roads[i][0]);
        }
        
        Queue<Integer> q = new LinkedList<>();
        q.add(destination);
        visited[destination] = true;
        dist[destination] = 0;
        
        while(!q.isEmpty()) {
            int now = q.poll();
            
            for(int next : road[now]) {
                if(!visited[next]) {
                    visited[next] = true;
                    q.add(next);
                    dist[next] = dist[now] + 1;
                }
            }
        }
        
        for(int i=0; i<sources.length; i++) {
            answer[i] = dist[sources[i]];
        }
        
        return answer;
    }
}

 

풀이

1. 처음엔 sources에서 destination으로 가는 것이기 때문에 sources를 for문을 돌면서 각각 bfs나 다익스트라 알고리즘을 이용하여 destination으로 가는 최단 거리를 구해서 풀었더니 시간 초과가 났다. 결국 sources에서는 모두 목적지가 destination으로 가려고 하는 것이기 때문에 destination으로 bfs나 다익스트라 알고리즘을 이용하여 한 번만 각 노드까지의 최단 거리를 1차원 배열로 구하여 destination에서 각 sources의 지점까지의 거리를 확인하면 된다. (이 문제의 경우 경로의 거리가 모두 1이기 때문에 굳이 다익스트라를 사용하지 않고 bfs를 이용해서 푼다. bfs 특성상 한 번 방문한 노드는 그 당시가 최단 거리이기 때문에 한 번 방문한 노드를 다시 방문하지 않도록 방문 체크 테이블(visited)을 만들고 방문할 때마다 체크하도록 한다.)
2. 우선 destination에서 각 노드까지의 거리를 구하기 위한 1차원 배열 dist와 방문 체크를 위한 visited를 생성하고 지역과 지역의 연결 그래프를 나타내기 위한 Integer 타입의 ArrayList 배열 road를 생성한다.
3. 이동이 불가능 한 경우에는 문제에 주어진대로 -1을 반환하기 위해 dist를 전부 -1로 초기화 한다.
4. 그리고 주어진 roads를 돌면서 배열 road의 해당 지역의 번호 인덱스에 해당하는 ArrayList에 연결 지역을 전부 추가하도록 한다. 
5. 이제 bfs 탐색을 하도록 한다.
6. Queue를 생성하고 destination을 담은 후에 destination의 방문체크를 위해 visited[destination]을 true로 변경하고 dist[destination]을 0으로 변경한다.
7. Queue가 빌 때까지 반복하는 while문을 생성하고 Queue의 첫 데이터를 뽑아 road에서 연결된 지역들을 for each문을 돌면서 연결된 지역이 방문하지 않았다면 해당 연결 지역의 visited를 true로 변경하고 Queue에 담고 dist[해당 연결 지역]을 dist[현재 지역] + 1을 하여 dist를 업데이트 하며 while문을 반복한다.
8. while문을 종료하면 dist에 destination에서 모든 지역까지의 최단 거리가 업데이트 되어 있고 갈 수 없는 지역은 처음에 초기화했던 -1로 되어으므로 sources를 for문을 돌면서 dist에서 해당 지역의 값을 꺼내 answer에 순서대로 저장한 후에 answer를 반환하면 된다.  

 

 

 

다익스트라 알고리즘 이용

import java.util.*;

// sources를 돌면서 다익스트라를 반복하는게 아니라 destination으로 다익스트라를 구한다.
class Solution {
    ArrayList<ArrayList<Integer>> graph;
    public int[] solution(int n, int[][] roads, int[] sources, int destination) {
        int[] answer = new int[sources.length];
		
        graph = new ArrayList<>();
		
        for(int i=0; i<=n; i++) {
            graph.add(new ArrayList<>());
        }
		
        for(int i=0; i<roads.length; i++) {
            graph.get(roads[i][0]).add(roads[i][1]);
            graph.get(roads[i][1]).add(roads[i][0]);
        }
		
        int[] dist = new int[n+1];

        Arrays.fill(dist, 100001);

        dist = dijkstra(destination, dist);
        
        for(int i=0; i<sources.length; i++) {
           if(dist[sources[i]] == 100001)
                answer[i] = -1;
           else
                answer[i] = dist[sources[i]];
        }
	
        return answer;
    }
   
    public int[] dijkstra(int start, int[] dist) {
        PriorityQueue<Node> pq = new PriorityQueue<>((o1, o2) -> o1.cost - o2.cost);
        pq.offer(new Node(start, 0));
        dist[start] = 0;
		
        while(!pq.isEmpty()) {
            Node now = pq.poll();
            int nIdx = now.idx;
            int nCost = now.cost; // 처음 start 지점에서 nIdx까지 거리
		   
            if(nCost > dist[nIdx]) continue;
		    
            for(int next : graph.get(nIdx)) {
                int cost = dist[nIdx] + 1;
		       
                if(cost < dist[next]) {
                    dist[next] = cost;
                    pq.offer(new Node(next, cost));
                }
            }
        }
        
        return dist;
    }
	
    class Node {
        int idx;
        int cost;
       
        public Node(int idx, int cost) {
            this.idx = idx;
            this.cost = cost;
        }
    }
}

 

풀이

1. 해당 풀이는 다익스트라 알고리즘을 이용한 풀이이다.
2. roads의 연결 지역을 나타내기 위한 각각의 원소에 ArrayList를 갖는 ArrayList인 graph를 생성한다.
3. graph의 인덱스가 해당 지역을 나타내기 위해 인덱스 n까지 해당 인덱스에 ArrayList를 생성한다.
4. 연결 상태를 나타내도록 주어진 roads를 돌면서 graph의 해당 지역의 인덱스에 있는 ArrayList에 연결 지역들을 모두 담는다.
5. destination으로부터 모든 지역의 최단 거리를 나타내기 위해 dist 배열을 생성한다.
6. dist 배열을 문제에서 나올 수 있는 최고의 수인 100000 보다 큰 수인 100001로 초기화 한다.
7. 이제 destination을 출발점으로 하여 다익스트라를 실행한다.
8. 해당 지역의 번호 idx와 지역까지의 거리 cost를 멤버 변수로 갖는 Node 클래스를 만든다.
9. Node를 선언 타입으로 하는 우선순위큐를 생성하고 cost를 기준으로 오름차순 정렬되도록 한다.
10. 우선순위큐에 idx를 출발점인 destination과 cost를 0으로 갖는 Node를 생성하여 담고 dist[출발점]을 0으로 변경한다.
11. 우선순위큐가 빌 때까지 반복하는 while문을 생성하고 우선순위큐에서 첫 데이터를 가져온다.
12. 해당 Node의 cost가 dist 해당지역보다 크다면 넘어가고(이미 cost가 더 크므로 기존의 dist가 최단 거리이므로) 아니라면 해당 지역에서 연결된 다음 지역들을 graph에서 가져와 for each문을 돌면서 연결된 지역의 cost를 임시로 현재 지역의 dist 값 + 1로 생각하여 해당 cost가 dist[연결된 다음 지역 번호] 보다 작다면 더 최단 거리이므로 dist[연결된 다음 지역 번호]를 해당 cost로 변경하고 우선순위큐에 idx를 연결된 다음 지역 번호와 해당 cost를 갖는 Node를 생성하여 담는다.
13. 이렇게 while문을 반복한 후 종료되면 dist에 출발점 destination으로부터 최단 거리가 업데이트 되어있고 갈 수 없는 지역은 기존에 최댓값으로 초기화 했던 100001으로 되어있다.
14. 따라서 sources를 for문을 돌면서 dist의 해당 값이 100001이라면 answer에 -1을 저장하고 아니라면 해당 dist 값을 저장하고 answer를 반환한다.