Kotlin에서의 다익스트라 구현

에제문제 : 프로그래머스(배달)

그래프는 어떻게 그리지?

In Python

graph = [[] for _ in range(n+1)]

for i in road:
	graph[i[0]].append((i[1],i[2]))
	graph[i[1]].append((i[0],i[2]))

In Kotlin

var graph = Array<ArrayList<Pair<Int,Int>>>(n+1){ ArrayList() }
for (i in road){
	graph[i[0]].add(Pair(i[1],i[2]))
	graph[i[1]].add(Pair(i[0],i[2]))
}

파이썬과는 다르게 2차원 Array안에 ArrayList를 생성해되, ArrayList는 Pair쌍을 가지도록

• Array
  • ArrayList
    • Pair

(n+1){ ArrayList() }은 Array의 파라미터 및 lamda식

//ex
var arr = Array(5){i->1}
println(arr.contentToString()) // [1, 1, 1, 1, 1]

우선순위큐는 어떻게?

In Python

dq = list()
heapq.heappush(dq,[3,2])
heapq.heappush(dq,[5,5])
heapq.heappush(dq,[1,4])
heapq.heappush(dq,[2,5])
print(dq)

In Kotlin

import java.util.PriorityQueue

private val dq = PriorityQueue<Pair<Int,Int>>({
    a,b->
    when{
        a.second < b.second -> -1
        a.second > b.second -> 1
        else -> 0
    }
})

fun main() {
    dq.add(Pair(3,2))
    dq.add(Pair(5,5))
    dq.add(Pair(1,4))
    dq.add(Pair(6,6))
    dq.add(Pair(2,5))
    dq.add(Pair(7,7))

    while(dq.isNotEmpty()){
        println(dq.poll())
    }

}

출력 값

(3, 2)
(1, 4)
(2, 5)
(5, 5)
(6, 6)
(7, 7)

Pair 객체는 뭐지?

다익스트라 알고리즘에서 (cost, destination)을 담는 배열이 필요하다.

배열에 배열을 포함하는 객체를 만들 수 있지만 배열을 포함한 클래스를 만드는 것은 피곤한 일이다.

그래서 코틀린은 Pair와 Triple 객체를 기본으로 제공

• Pair는 2개의 객체를 저장할 수 있는 객체
• Triple은 3개의 객체를 저장하는 객체

val (a, b) = Pair(1, "x")
println(a) // 1
println(b) // x

fun <T> Pair<T, T>.toList(): List<T>
(source)

우선순위 큐는 이에 Comparator을 사용하여 이를 정렬한다.

ex)
println(listOf("aaa", "bb", "c").sortedWith(compareBy { it.length })) 
// output: // [c, bb, aaa]

출처: https://iosroid.tistory.com/94

private val dq = PriorityQueue<Pair<Int,Int>>(Comparator{
    a,b->
    when{
        a.second < b.second -> -1
        a.second > b.second -> 1
        else -> 0
    }
})

//val (cur,curDis) = pq.poll()

전체소스

전체소스 Python

import heapq

def solution(n, road, k):
    graph = [[] for _ in range(n+1)]

    for i in road:
        graph[i[0]].append((i[1],i[2]))
        graph[i[1]].append((i[0],i[2]))

    dist = [10**9] * (n+1)
    def dijkstra(start):
        dist[start] = 0
        q = list()
        heapq.heappush(q, [0,start])

        while q:
            curdist, now = heapq.heappop(q)#0, 1
            if dist[now] < curdist:
                continue
            for i in graph[now]:
                cost = curdist + i[1]
                if cost < dist[i[0]]:
                    dist[i[0]] = cost
                    heapq.heappush(q,[cost,i[0]])

    dijkstra(1)
    cnt =0
    for i in dist:
        if i <=k:
            cnt+=1
    return cnt

전체소스 Kotlin

import java.util.PriorityQueue
class Solution {
    private val dq = PriorityQueue<Pair<Int,Int>>(Comparator{
        a,b->
        when{
            a.second < b.second -> -1
            a.second > b.second -> 1
            else -> 0
        }
    })

    fun solution(n: Int, road: Array<IntArray>, k: Int): Int {
        var graph = Array<ArrayList<Pair<Int,Int>>>(n+1){ ArrayList() }
        for (i in road){
            graph[i[0]].add(Pair(i[1],i[2]))
            graph[i[1]].add(Pair(i[0],i[2]))
        }

        fun dijkstra(start:Int): IntArray {
            var distances = IntArray(n+1,{i->987654321})
            distances[start] = 0
            dq.add(Pair(start,0))
            while(dq.isNotEmpty()){
                var (now,nowdist) = dq.poll()
                if (distances[now] < nowdist) continue
                for(i in graph[now]){
                    var cost = nowdist + i.second
                    var nextdes = i.first
                    if (distances[nextdes] > cost){
                        distances[nextdes] = cost
                        dq.add(Pair(nextdes,cost))
                    }
                }
            }
            return distances
        }

        var ans = dijkstra(1)
        var cnt = 0
        for(i in ans){
            println(i)
            if (i<=k){
                cnt+=1
            }
        }
        return cnt
    }
}