라떼는말이야

코틀린 다익스트라 알고리즘 Kotlin dijkstra algorithm 본문

알고리즘/코딩 테스트

코틀린 다익스트라 알고리즘 Kotlin dijkstra algorithm

MangBaam 2022. 8. 16. 03:33
반응형

이 글의 대상

기본적인 다익스트라 알고리즘에 대해 알고 있다고 가정하고 코틀린을 사용해서는 어떻게 다익스트라 알고리즘을 구현할 지 궁금한 사람들을 대상으로 하고 있다

 


0. 입출력 형태

문제마다 입력 형태나 출력 형태가 각이각색이다. 이번 글에서 소개하는 코드는 다음 형태로 작성되었다.

  1. 노드 개수, 간선 개수 입력 받기
  2. 시작 노드 입력 받기
  3. 간선 정보 입력 받기
  4. 2번에서 입력 받은 시작 노드부터 탐색 시작
  5. 노드 번호 순으로 시작 노드로부터 거리 출력 (갈 수 없는 노드는 INFINITE 출력)

 

1. 노드 개수, 간선 개수 입력 받기

import java.util.*

fun dijkstra() {
    val (n, m) = readln().trim().split(" ").map { it.toInt() } // 노드 개수, 간선 개수
}

 

2. 시작 노드 입력 받기

import java.util.*

fun dijkstra() {
    val (n, m) = readln().trim().split(" ").map { it.toInt() } // 노드 개수, 간선 개수
    val start = readln().toInt() // 시작 노드
    val graph = Array(n + 1) { mutableListOf<Pair<Int, Int>>() } // 노드 연결 정보
    val distance = (0..n).map { Int.MAX_VALUE }.toMutableList() // 최단 거리 테이블

}

1번과 2번에서 입력 받은 정보 외에 graph 와 distance 라는 것도 생성했다. graph 는 노드 간 연결 정보를 담는 자료구조이다. 배열의 각 아이템은 MutableList 로 구성되어 있고, 그 안에는 Pair 가 들어간다.

Pair 는 2 개의 데이터를 가지는 자료 구조이다. 그리고 각각은 Pair.first, Pair.second 로 접근할 수 있다.

Pair 의 첫 번째 요소는 노드 번호이고, 두 번째 요소는 거리이다.

만약 위와 같이 연결된 그래프가 있을 때 graph[1] 에는 MutableList(Pair(2, 2), Pair(3, 5), Pair(4, 1)) 이 들어있을 것이다.

 

distance 는 시작 노드부터 각 노드까지의 거리가 저장된다. 초기 값은 Int.MAX_VALUE 로 초기화된다. (물론 문제에서 노드 간 거리가 주어질 때 Int.MAX_VALUE 보다 큰 값이 주어지는지 확인이 필요하다)

 

3. 간선 정보 입력 받기

import java.util.*

fun dijkstra() {
    val (n, m) = readln().trim().split(" ").map { it.toInt() } // 노드 개수, 간선 개수
    val start = readln().toInt() // 시작 노드
    val graph = Array(n + 1) { mutableListOf<Pair<Int, Int>>() } // 노드 연결 정보
    val distance = (0..n).map { Int.MAX_VALUE }.toMutableList() // 최단 거리 테이블

    /* 모든 간선 입력 */
    repeat(m) {
        val (a, b, c) = readln().split(" ").map { it.toInt() }
        graph[a].add(Pair(b, c)) // a -> b 비용: c
    }
}

m 개의 간선이 있다고 했으니 m 번 반복해서 graph 에 저장한다.

1 3 5 가 입력되면 1번 노드에서 3번 노드로 가는 거리가 5라는 뜻이다.

 

4. 2번에서 입력 받은 시작 노드부터 탐색 시작 (dijkstraSearch 함수는 뒤에서 설명)

import java.util.*

fun dijkstra() {
    val (n, m) = readln().trim().split(" ").map { it.toInt() } // 노드 개수, 간선 개수
    val start = readln().toInt() // 시작 노드
    val graph = Array(n + 1) { mutableListOf<Pair<Int, Int>>() } // 노드 연결 정보
    val distance = (0..n).map { Int.MAX_VALUE }.toMutableList() // 최단 거리 테이블

    /* 모든 간선 입력 */
    repeat(m) {
        val (a, b, c) = readln().split(" ").map { it.toInt() }
        graph[a].add(Pair(b, c)) // a -> b 비용: c
    }

    dijkstraSearch(start) // 탐색 수행
}

dijkstraSearch 함수는 뒤에서 설명할 예정

 

5. 노드 번호 순으로 시작 노드로부터 거리 출력 (갈 수 없는 노드는 INFINITE 출력)

import java.util.*

fun dijkstra() {
    val (n, m) = readln().trim().split(" ").map { it.toInt() } // 노드 개수, 간선 개수
    val start = readln().toInt() // 시작 노드
    val graph = Array(n + 1) { mutableListOf<Pair<Int, Int>>() } // 노드 연결 정보
    val distance = (0..n).map { Int.MAX_VALUE }.toMutableList() // 최단 거리 테이블

    /* 모든 간선 입력 */
    repeat(m) {
        val (a, b, c) = readln().split(" ").map { it.toInt() }
        graph[a].add(Pair(b, c)) // a -> b 비용: c
    }

    dijkstraSearch(start) // 탐색 수행

    distance.drop(1).forEach { d ->
        if (d == Int.MAX_VALUE) println("INFINITY") else println(d)
    }
}

다익스트라 탐색이 완료되면 distance 에는 시작 노드부터 각 노드까지의 최단 거리가 저장되어 있다.

순차적으로 출력하되, distance 의 0번 인덱스는 사용하지 않았기 때문에 drop(1) 로 제외하고 탐색한다.

거리가 Int.MAX_VALUE 와 같다면 탐색할 수 없는 노드인 것이기에 INFINITY 를 출력한다.

 

다익스트라 최단 거리 탐색

다익스트라 최단 거리 탐색 알고리즘에서 큐에서 가장 거리가 짧은 것을 찾아야 하는 과정이 필요한데 이 탐색 과정을 for 문으로 선형 탐색하게 되면 O(N) 의 시간이 소요된다. 이 부분을 우선 순위 큐로 구현하면 시간 복잡도를 훨씬 줄일 수 있다.

import java.util.*

fun dijkstra() {
    // 정보 입력

    /* 모든 간선 입력 */
    // ...

    /* 다익스트라 최단 거리 탐색 */
    fun dijkstraSearch(start: Int) {
        val q = PriorityQueue<Pair<Int, Int>> { p1, p2 -> // 우선 순위 큐
            when {
                p1.second > p2.second -> 1
                p1.second < p2.second -> -1
                else -> 0
            }
        }

        q.add(Pair(start, 0))
        distance[start] = 0

        while (q.isNotEmpty()) {
            val (now, dist) = q.poll() // 최단 거리가 가장 짧은 정보 꺼내기
            if (distance[now] < dist) continue // 이미 처리된 적 있는 노드

            /* 현재 노드와 연결된 인접 노드 확인 */
            for (i in graph[now]) {
                val cost = dist + i.second
                if (cost < distance[i.first]) { // 현재 노드를 거쳐서 다른 노드로 이동하는 거리가 더 짧은 경우
                    distance[i.first] = cost
                    q.add(Pair(i.first, cost))
                }
            }
        }
    }

    // 탐색 수행 ...

    // 최단 거리 출력 ...
}

코틀린에서는 PriorityQueue 라는 자료 구조가 제공된다. 이 것을 사용하면 쉽게 우선 순위 큐를 구현할 수 있다. 생성자의 마지막 파라미터로는 고차 함수가 주어지는데 이 람다식에서 우선 순위를 지정할 수 있다. 이 우선 순위는 숫자로 표현되는데 주로 -1, 0, 1 로 표현된다.

위 코드의 식은 최단 거리를 찾기 위한 조건이고, 만약 부등호를 뒤집으면 최장 거리를 찾게 될 것이다.

나머지 부분은 다익스트라 알고리즘을 알고 있다면 무슨 의미인지 보일 것이다.

 

전체 코드

import java.util.*

fun dijkstra() {
    val (n, m) = readln().trim().split(" ").map { it.toInt() } // 노드 개수, 간선 개수
    val start = readln().toInt() // 시작 노드
    val graph = Array(n + 1) { mutableListOf<Pair<Int, Int>>() } // 노드 연결 정보
    val distance = (0..n).map { Int.MAX_VALUE }.toMutableList() // 최단 거리 테이블

    /* 모든 간선 입력 */
    repeat(m) {
        val (a, b, c) = readln().split(" ").map { it.toInt() }
        graph[a].add(Pair(b, c)) // a -> b 비용: c
    }

    /* 다익스트라 최단 거리 탐색 */
    fun dijkstraSearch(start: Int) {
        val q = PriorityQueue<Pair<Int, Int>> { p1, p2 -> // 우선 순위 큐
            when {
                p1.second > p2.second -> 1
                p1.second < p2.second -> -1
                else -> 0
            }
        }

        q.add(Pair(start, 0))
        distance[start] = 0

        while (q.isNotEmpty()) {
            val (now, dist) = q.poll() // 최단 거리가 가장 짧은 정보 꺼내기
            if (distance[now] < dist) continue // 이미 처리된 적 있는 노드

            /* 현재 노드와 연결된 인접 노드 확인 */
            for (i in graph[now]) {
                val cost = dist + i.second
                if (cost < distance[i.first]) { // 현재 노드를 거쳐서 다른 노드로 이동하는 거리가 더 짧은 경우
                    distance[i.first] = cost
                    q.add(Pair(i.first, cost))
                }
            }
        }
    }

    dijkstraSearch(start) // 탐색 수행

    distance.drop(1).forEach { d ->
        if (d == Int.MAX_VALUE) println("INFINITY") else println(d)
    }
}

 

시간 복잡도

위 코드의 시간 복잡도는 O(ElogV) 이다. (E: 간선의 개수, V: 노드의 개수)

우선 우선 순위 큐는 Heap 구조로 되어있기 때문에 O(logN) 이다.

다익스트라 알고리즘을 전체적으로 보면 각 노드 별로 인접한 노드의 최단 거리를 확인하는 작업을 반복한다. 즉, 간선의 개수 E 개 만큼 우선 순위 큐에 넣었다가 모두 빼는 과정을 한다고 생각할 수 있다.

이렇게 생각했을 때 O(ElogE) 의 복잡도를 가지게 된다.

 

E 는 항상 V^2 보다 작거나 같다. 모든 노드가 서로 연결되었을 때 최대 V^2 개의 간선이 생길 수 있기 때문이다.

그래서 O(ElogE) <= O(ElogV^2) 이고, O(ElogV^2) 는 정리하면 O(2ElogV) -> O(ElogV) 가 된다.

결국 O(ElogE) <= O(ElogV) 가 되고, 시간 복잡도를 O(ElogV) 라고 할 수 있다.

 

반응형
Comments