Dijkstra PriorityQueue 사용해서 구현해보기 ([백준] 1753번)

https://www.acmicpc.net/problem/1753

1753번: 최단경로

문제

방향그래프가 주어지면 주어진 시작점에서 다른 모든 정점으로의 최단 경로를 구하는 프로그램을 작성하시오. 단, 모든 간선의 가중치는 10 이하의 자연수이다.

입력

첫째 줄에 정점의 개수 V와 간선의 개수 E가 주어진다. (1≤V≤20,000, 1≤E≤300,000) 모든 정점에는 1부터 V까지 번호가 매겨져 있다고 가정한다. 둘째 줄에는 시작 정점의 번호 K(1≤K≤V)가 주어진다. 셋째 줄부터 E개의 줄에 걸쳐 각 간선을 나타내는 세 개의 정수 (u, v, w)가 순서대로 주어진다. 이는 u에서 v로 가는 가중치 w인 간선이 존재한다는 뜻이다. u와 v는 서로 다르며 w는 10 이하의 자연수이다. 서로 다른 두 정점 사이에 여러 개의 간선이 존재할 수도 있음에 유의한다.

출력

첫째 줄부터 V개의 줄에 걸쳐, i번째 줄에 i번 정점으로의 최단 경로의 경로값을 출력한다. 시작점 자신은 0으로 출력하고, 경로가 존재하지 않는 경우에는 INF를 출력하면 된다.

---

최단거리를 구할 때 가장 대표적으로 사용하는 Dijkstra 알고리즘을 사용하여 푸는 문제이다. 

Dijkstra 알고리즘을 구현하는 방법은 크게 두 가지가 있는데, 하나는 다익스트라 알고리즘 정의 그대로 구현하는 것이고 다른 하나의 방법은 Heap을 사용하는 방법이다. Heap을 사용하면 더 빠른 시간복잡도를 사용하여 문제를 해결할 수 있기 때문에 주로 Dijkstra를 구현할 때는 heap을 사용하는 것을 선호하는 편이다.

 

Java에서 Heap을 구현하기 위해 우리는 PriorityQueue를 사용할 것인데 자바에서 PriorityQueue는 min heap이기 때문에 poll을 하면 가장 작은 원소가 나온다. 따라서 다익스트라를 구현할 때 별다른 조치 없이 그대로 사용해도 무방하다. 

 

---

먼저 위 문제는 weight가 있는 그래프이기 때문에 Node라는 class를 하나 더 새롭게 정의할 것이고, 여기서 주의해야 할 점은 PriorityQueue에 사용될 자료형이 반드시 Comparable 구현이 되어 있어야 한다는 점이다. 새롭게 만들어준 class이기 때문에 크기/순서를 직접 정의해주지 않는다면 Java도 어떤 식으로 heap이 구성되어야 하는지 혼란이 올 것이기 때문이다. 

class Node implements Comparable<Node>{
    int index;
    int weight;
    Node(int index, int weight){
        this.index=index;
        this.weight=weight;
    }
    public int compareTo(Node o) {
        return Integer.compare(this.weight, o.weight);
    }
}

따라서 위와 같이 Comparable을 implements해주고 객체 크기의 기준은 weight로 비교하도록 했다. 여기서 weight란 각 node에 누적된 distance라고 봐도 무방하다. 

---

큰 logic을 설명하자면 다음과 같다. (여기서 다익스트라는 방문한 node를 또 방문하며 갱신해줄 수 있기 때문에 방문을 했는지 여부를 나타내는 배열은 굳이 필요하지 않다)

1. d[]라는 배열에는 각 node index의 계산된 최단거리를 계산한다.

2. d 배열을 무한대로 초기화하고, 시작 index의 d 값만 0으로 세팅한다.

3. PriorityQueue에 처음 시작 index를 넣는다.

4. queue가 empty가 되지 않을 때까지 while 문을 돌리며 하나씩 넣고, 인접한 것들 중 값이 더 작아지는 node를 넣는 과정을 진행한다. 

5. 여기서 갱신이란, 인접한 node의 d값이 queue에서 꺼낸 node의 d 값 + 간선의 weight보다 클 때 일어나는 현상이다. 더 짧은 path가 존재하므로 갱신해준다는 느낌이다.

6. 이 과정이 끝나면 각 vertex를 돌면서 출력을 해주면 된다. 

 

이 모든 과정을 코드로 설명하면 아래와 같다. 

이 문제가 아니라 다익스트라 코드가 궁금한 사람들은 아래 주석에 다익스트라 부분이라고 써놓았으니 그 부분을 체크해보면 될 것 같다. 

import java.util.*;
class Node implements Comparable<Node>{
    int index;
    int weight;
    Node(int index, int weight){
        this.index=index;
        this.weight=weight;
    }
    public int compareTo(Node o) {
        return Integer.compare(this.weight, o.weight);
    }
}
public class Main{
    public static void main(String[] args){
        Scanner sc=new Scanner(System.in);
        int v=sc.nextInt();
        int e=sc.nextInt();
        int start=sc.nextInt();
        ArrayList<Node> arr[]=new ArrayList[v+1];
      
        for(int i=1; i<=v; i++){
            arr[i]=new ArrayList<>();
        }
 
        for(int i=0; i<e; i++){
            int s=sc.nextInt();
            int en=sc.nextInt();
            int weight=sc.nextInt();
            arr[s].add(new Node(en, weight));
        }
        
        
        //다익스트라 부분 시작 
       PriorityQueue<Node> pq=new PriorityQueue<>();
       int d[]=new int[v+1];
        Arrays.fill(d, Integer.MAX_VALUE);
        d[start]=0;
        pq.add(new Node(start, 0));
        while(!pq.isEmpty()){
            Node current=pq.poll();
            for(Node next: arr[current.index]){
                if (d[next.index]>d[current.index]+next.weight){
                    d[next.index]=d[current.index]+next.weight;
                    pq.add(new Node(next.index, d[next.index]));
                }
            }
        }
    //다익스트라 부분 끝
      StringBuilder stringBuilder = new StringBuilder();
        for(int i=1; i<=v; i++){
            if(d[i]!=Integer.MAX_VALUE){
                    stringBuilder.append(d[i] + "\n");
            }
            else{
                 stringBuilder.append("INF\n");
            }  
        }
        System.out.print(stringBuilder.toString());
      
    }
         
}

 

참고로 다익스트라라는 알고리즘은 간선의 길이가 non negative 일때, 즉 음이 아닐 때만 사용이 가능하다. 

음수 간선이 있을 때는 Bellman-Ford 알고리즘을 사용한다. 

기회가 있으면 Bellman-Ford 알고리즘에 대해서도 포스팅해보도록 하겠다.