[알고리즘] Dijkstra 알고리즘 - PriorityQueue

[알고리즘] 다익스트라(Dijkstra) 알고리즘

이전 포스팅의 다익스트라 알고리즘을 구현한 코드는 O(N^2)의 시간 복잡도를 갖는다. 모든 정점을 한 번씩 순회하는 N번의 loop와 그 안에서 distance[i]의 가장 작은 값의 인덱스를 찾는 N번의 loop가 호출되었다. 단순히 다익스트라 알고리즘을 구현하고 사용하는데에는 코드도 직관적이고 다른 자료구조도 안쓰여서 쉽게 쓸 수 있으나 문제를 풀거나, 서비스에 적용해야 할 때는 문제가 될 여지가 충분하다.

 

 이를 보완하고자 우선 순위 큐(Priority Queue)를 사용하고 인접리스트를 만들어 다익스트라 알고리즘을 구현하면 시간 복잡도를 O(Elog V) 까지 줄일 수 있다. 만약 인접 행렬로 구현했다면 매번 인접한 정점을 찾아야 하므로 루프마다 O(V)의 시간이 소요되어 O(V^2)가 될 것이다.

---

    static class Edge implements Comparable<Edge> {
        int target, weight;
        public Edge(int t, int w){
            target = t;
            weight = w;
        }

        @Override
        public int compareTo(Edge edge) {
            return this.weight - edge.weight;
        }
    }

간선의 정보를 담을 Edge 클래스이다. Edge 클래스는 target 정점과 가중치 weight를 갖으며, PriorityQueue에 들어가기 위해선 비교할 수 있는 기준이 필요하므로, Comparable 인터페이스를 구현해 우선 순위 기준을 마련했다.

    static void dijkstra(int start){

        distance[start] = 0;

        Q.add(new Edge(start, distance[start]));

        while(!Q.isEmpty()){
            Edge edge = Q.poll(); //우선순위 큐에서 꺼냄

            int current = edge.target; //현재 방문하고 있는 정점
            int weight = edge.weight; //그 가중치

            if(distance[current] < weight) continue; //갱신할 필요가 없으면 skip

            for(int i=0; i<list[current].size(); i++){ //현재 방문한 정점의 인접 정점 iteration
                Edge nextEdge = list[current].get(i);

                int next = nextEdge.target;
                int nextWeight = weight + nextEdge.weight; //갱신할 값

                if(nextWeight < distance[next]){

                    distance[next] = nextWeight;
                    Q.add(new Edge(next, nextWeight));
                }
            }
        }
    }

다익스트라 알고리즘의 소스 코드이다. start정점의 인접리스트에서 원소를 꺼내어 큐에 넣고 while문을 돌린다.

우선순위 큐에서 하나씩 꺼내어 최소 가중치(distance[])를 더 적은 값으로 갱신할 수 있으면 갱신하도록 한다.

# Code </>

public class DijkstraPriorityQ {

    static int n = 6;

    static class Edge implements Comparable<Edge> {
        int target, weight;
        public Edge(int t, int w){
            target = t;
            weight = w;
        }

        @Override
        public int compareTo(Edge edge) {
            return this.weight - edge.weight;
        }
    }

    static Queue<Edge> Q = new PriorityQueue<>();
    static List<Edge>[] list = new ArrayList[n];

    static final int INF = 1000000000;

    static int[] distance = new int[n];

    static void dijkstra(int start){

        distance[start] = 0;

        Q.add(new Edge(start, distance[start]));

        while(!Q.isEmpty()){
            Edge edge = Q.poll(); //우선순위 큐에서 꺼냄

            int current = edge.target; //현재 방문하고 있는 정점
            int weight = edge.weight; //그 가중치

            if(distance[current] < weight) continue; //갱신할 필요가 없으면 skip

            for(int i=0; i<list[current].size(); i++){ //현재 방문한 정점의 인접 정점 iteration
                Edge nextEdge = list[current].get(i);

                int next = nextEdge.target;
                int nextWeight = weight + nextEdge.weight; //갱신할 값

                if(nextWeight < distance[next]){

                    distance[next] = nextWeight;
                    Q.add(new Edge(next, nextWeight));
                }
            }
        }
    }

    public static void main(String[] args) {

        /*
        int[][] arr = {
                {0, 7, 9, INF, INF, 14},
                {7, 0, 10, 15, INF, INF},
                {9, 10, 0, 11, INF, 2},
                {INF, 15, 11, 0, 6, INF},
                {INF, INF, INF, 6, 0, 9},
                {14, INF, 2, INF, 9, 0}
        };
         */

        for(int i=0; i<n; i++)
            distance[i] = INF;

        for(int i=0; i<n; i++)
            list[i] = new ArrayList<>();

        //인접 리스트에 모두 할당
        list[0].add(new Edge(5, 14));
        list[0].add(new Edge(1, 7));
        list[0].add(new Edge(2, 9));


        list[1].add(new Edge(0, 7));
        list[1].add(new Edge(2, 10));
        list[1].add(new Edge(3, 15));

        list[2].add(new Edge(0, 9));
        list[2].add(new Edge(1, 10));
        list[2].add(new Edge(3, 11));
        list[2].add(new Edge(5, 2));

        list[3].add(new Edge(1, 15));
        list[3].add(new Edge(2, 11));
        list[3].add(new Edge(4, 6));

        list[4].add(new Edge(3, 6));
        list[4].add(new Edge(5, 9));

        list[5].add(new Edge(0, 14));
        list[5].add(new Edge(2, 2));
        list[5].add(new Edge(4, 9));

        dijkstra(0);

        Arrays.stream(distance).forEach(i->System.out.print(i + " "));
    }
}

 

# References

동빈나님의 유튜브

안경잡이 개발자님의 블로그

Ries 마법의 슈퍼마리오님의 블로그