[알고리즘][KOTLIN] 다익스트라(w. 우선순위큐, 인접 리스트)

 

[알고리즘][KOTLIN] 다익스트라

위의 알고리즘을 풀 때 노드 및 간선의 입력이 커지게 되면 기본 다익스트라 알고리즘의 경우 O(n^2) 시간 복잡도를 가지므로 인접 리스트를 이용하여 시간복잡도를 해결해야하는 상황이 발생할 수 있다.

코드

import java.io.*
import java.util.*

fun main(args: Array<String>) {
    val br = BufferedReader(InputStreamReader(System.`in`))
    // V : 노드의 개수
    // E : 간선의 개수
    val (V, E) = br.readLine().split(" ").map { it.toInt() }
    val start = br.readLine().toInt() // 시작 노드

    val g = Graph(V)

	// 간선 정보를 입력한다.
    repeat(E) {
        val (s, e, w) = br.readLine().split(" ").map { it.toInt() }
        g.input(s, e, w)
    }

	// 다익스트라 알고리즘을 수행한다.
    g.dijkstra(start)
}


private class Graph(private val n: Int) {
    private val list = Array(n + 1) {
        mutableListOf<Node>()
    }

    fun input(start: Int, end: Int, weight: Int) {
        list[start].add(Node(end, weight))
    }

    fun dijkstra(start: Int) {
        val check = BooleanArray(n + 1)
        val distance = IntArray(n + 1)
        Arrays.fill(distance, Int.MAX_VALUE)
        distance[start] = 0

		// 우선 순위 큐에 현재 노드 정보를 저장한다.
        val priorityQueue = PriorityQueue<Node>()
        priorityQueue.offer(Node(start, 0))

        while (!priorityQueue.isEmpty()) {
            val current = priorityQueue.poll().idx

            if (check[current]) continue

            check[current] = true

            for (next in list[current]) {
                if (next.weight != 0 && next.weight + distance[current] < distance[next.idx]) {
                    distance[next.idx] = next.weight + distance[current]
                    priorityQueue.offer(Node(next.idx, distance[next.idx]))
                }
            }
        }

        for (i in 1..n) {
            if (distance[i] == Int.MAX_VALUE) println("INF")
            else println(distance[i])
        }
    }
}

// 인접 리스트를 위한 노드 정보를 저장하고 비교할 수 있는 Class
private class Node(val idx: Int, val weight: Int) : Comparable<Node> {
    override fun compareTo(other: Node): Int {
        return weight - other.weight
    }
}