Today I Learned

오늘은 저번에 정리했던 다익스트라 알고리즘을 실제로 적용할 수 있는 문제를 찾아서 풀어봤다.

배달

문제

N개의 마을로 이루어진 나라가 있다. 1~N까지 번호가 붙여져 있다. 마을끼리는 양방향으로 이어진 도로가 있는데 int[][] road로 소요 시간이 {start,end,time}이 주어진다. 1번 마을에서부터 K시간안에 배달할 수 있는 마을의 수를 return 해라.

풀이 과정

1. 최소 가중치를 찾는 가장 기본적인 형태의 다익스트라 알고리즘 문제이다.

2. 우선 pq에 넣어야될 Node 클래스를 새로 생성한다. pq 내에서 비교할 수 있게 Comparable 인터페이스를 구현하고 compareTo 메서드를 override한다.

3. 인접리스트 형태로 ArrayList를 써서 그래프를 구현한다. 양방향그래프이므로 start와 end 양 지점에서 다 연결한다. 위에서 구현한 클래스인 Node를 사용한다.
   그리고 문제에서는 1~N개의 마을이므로 index를 맞춰주기위해 0~N으로 그래프를 만들고 index 0은 그냥 없는 셈 친다.

4. 우선순위큐를 생성하고, 각 마을까지의 최소시간을 저장할 배열을 생성한 뒤 index 1번만 제외하고 최대값(Integer.MAX_VALUE)을 넣는다. (0번은 갈일이 없으므로 최대값으로 둔다.)

5. pq에 자기자신을 넣고 while문으로 pq가 텅 빌때까지 아래의 반복문을 실행한다.
   5-1. pq에서 최소값의 노드를 꺼낸다.
   5-2. 해당 노드의 가중치가 일단 최소시간 배열의 가중치보다 작으면 일단 pass
   5-3. 그래프 상에서 해당 노드에 연결된 노드를 순회하면서 최소시간보다 nw+lw가 작으면 최소시간 배열에서 이 값으로 갱신하고 이를 다시 pq에 넣는다.

6. 위의 과정으로 "1번마을부터 다른 마을까지의 배달 최소 시간" 배열이 leastTime에 완성된다. 이 배열을 순회하면서 K값 이하 마을들을 count해 반환하면 끝.

구현 코드

import java.util.*;
class Solution {
    List<List<Node>> graph;
    int[] leastTime;
    int n;
    
    class Node implements Comparable<Node> {
        int vertex;
        int weight;
        
        Node(int vertex, int weight){
            this.vertex = vertex;
            this.weight = weight;
        }
        
        @Override
        public int compareTo(Node other){
            return Integer.compare(this.weight, other.weight);
        } 
    }
    
    void dijkstra(int index){
        Queue<Node> pq = new PriorityQueue<>();
        
        // 편의상 0~N까지로 배열을 만들고 0은 없는셈 친다.
        leastTime = new int[n+1];
        Arrays.fill(leastTime, Integer.MAX_VALUE);
        
        leastTime[1] = 0;
        pq.offer(new Node(index, 0));
        
        // pq에 담긴 최소시간의 node를 꺼내서 해당 노드와 연결된 간선들을 순회하며 최소 시간을 갱신한다.
        while(!pq.isEmpty()){
            Node node = pq.poll();
            int nv = node.vertex;
            int nw = node.weight;
            
            if(nw>leastTime[nv]) continue;
            
            for(Node linkedNode : graph.get(nv)){
                int lv = linkedNode.vertex;
                int lw = linkedNode.weight;
                
                if(nw+lw < leastTime[lv]){
                    leastTime[lv] = nw+lw;
                    pq.offer(new Node(lv, leastTime[lv]));
                }
            }
        }
        
    }
    
    
    
    public int solution(int N, int[][] road, int K) {
        int answer = 0;
        n = N;
        
        // 그래프 생성 0~N
        graph = new ArrayList<>();
        for(int i=0; i<=N; i++){
            graph.add(new ArrayList<>());
        }
        
        // 그래프 연결
        for(int[] r : road){
            int start = r[0];
            int end = r[1];
            int time = r[2];
            
            graph.get(start).add(new Node(end, time));
            graph.get(end).add(new Node(start, time));
        }
        
        // 1번 마을에서 시작
        dijkstra(1);
        
        // K 시간 이하의 마을 개수를 count
        for(int time : leastTime){
            if(time <= K) answer++;
        }
        
        return answer;
    }
}