Today I Learned

오늘도 코딩테스트 연습을 했다. 한동안 안써서 가물가물했던 다익스트라 알고리즘 문제를 풀어봤다.

부대 복귀

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

문제

부대원들이 여러 지역에 흩어져 임무 수행을 마치고 부대로 최단 시간에 복귀하려고 한다. 총 지역 수 n, 두 지역간 길을 담은 2차원 int배열 roads, 각 부대원이 위치한 서로 다른 지역 sources, 부대의 위치 destination이 주어진다. 참고로 두 지역간 길을 통과하는 데 걸리는 시간은 1로 동일하다. 주어진 조건에서 각 부대원들의 복귀 최단 시간을 배열에 담아 return해라. 길이 없어 복귀를 할 수 없는 부대원은 -1로 표기해라.

제한사항
3 ≤ n ≤ 100,000
2 ≤ roads의 길이 ≤ 500,000 (동일한 정보가 중복으로 주어지지 않는다.)
1 ≤ sources의 길이 ≤ 500
1 ≤ destination ≤ n

풀이 과정

전형적인 다익스트라 알고리즘 문제이다. 그래프를 구현하고 우선순위큐를 이용해 다익스트라 알고리즘을 구현하면 된다.

1. Node 클래스를 Comparable을 구현해 생성한다. 각 지역의 위치 v와 길이(시간, 가중치) w를 Node의 필드에 담고, w를 통해 비교하도록 구현한다.

2. roads의 정보를 이용해 이중리스트로 graph를 구현한다. 양방향으로 구현해주되, 길은 무조건 시간이 1이 소모되므로 1을 넣어준다.

3. 다익스트라 알고리즘 메서드를 우선순위큐를 이용해 구현해준다. 주어진 문제 조건 상 가장 큰 값은 100000이므로 max를 final로 100000으로 선언해두고 distance에다 채워둔다. 인덱스와 실제 지역번호가 맞지 않을 수 있으므로 N=n+1로 만들어 매칭시켜두고 0번 인덱스는 활용하지 않는 방식으로 한다.

4. 다익스트라 메서드를 destination을 출발지로 해서 실행시키면 distance 배열에 각 지역의 거리(시간)이 반환되어 나온다. distance 배열과 sources배열을 이용해 각 부대원이 부대 복귀하는 시간을 answer배열에 넣는다. 시간이 max면 -1, 아니면 그대로 넣으면 된다. answer 배열을 반환하면 완료!

Java 코드

import java.util.*;
class Solution {
    int[] distance;
    List<List<Node>> graph;
    final int max = 100000;
    int N;
    
    
    public class Node implements Comparable<Node> {
        int v;
        int w;
        
        public Node(int v, int w){
            this.v = v;
            this.w = w;
        }
        
        @Override
        public int compareTo(Node other){
            return Integer.compare(this.w, other.w);
        }
    }
    
    public void dijkstra(int index){
        Queue<Node> pq = new PriorityQueue<>();
        distance = new int[N];
        Arrays.fill(distance, max);
        distance[index] = 0;
        pq.offer(new Node(index, 0));
        
        while(!pq.isEmpty()){
            Node node = pq.poll();
            int nv = node.v;
            int nw = node.w;
            
            if (nw > distance[nv]) continue;
            
            for (Node linked : graph.get(nv)){
                int lv = linked.v;
                int lw = linked.w;
                if(nw+lw < distance[lv]){
                    distance[lv] = nw+lw;
                    pq.offer(new Node(lv, distance[lv]));
                }
            }
        }
    }
    
    
    
    public int[] solution(int n, int[][] roads, int[] sources, int destination) {
        N = n+1;
        
        graph = new ArrayList<>();
        for(int i=0; i<N; i++){
            graph.add(new ArrayList<>());
        }
        
        for(int[] road : roads){
            graph.get(road[0]).add(new Node(road[1],1));
            graph.get(road[1]).add(new Node(road[0],1));
        }
        
        dijkstra(destination);
        
        int[] answer = new int[sources.length];
        for(int i=0; i<sources.length; i++){
            int ans = distance[sources[i]];
            if(ans==max) answer[i] = -1;
            else answer[i] = ans;
        }
        
        
        return answer;
    }
}