코틀린(kotlin) - 우선순위큐(PriorityQueue) 정렬 기준

백준 또는 CodeForce 등을 풀다보면 우선순위 큐를 사용할 일이 많다.

이 때 우선 순위 큐에 들어가는 자료형이 Int, Double, String 등 한개로만 이루어져 있다면, 아무 조건없이 넣었을 때 오름차순으로 정렬될테고, reverseOrder() 명령어를 통해 역순으로 정렬할 수도 있다.

val pq = PriorityQueue<Int>()
pq.add(1)
pq.add(2)
pq.add(3)

val rpq = PriorityQueue<Int>(reverseOrder())
rpq.add(1)
rpq.add(2)
rpq.add(3)

/*
pq -> 1, 2, 3
rpq -> 3, 2, 1
*/

하지만 자료형이 하나가 아닌, Pair 또는 Triple 혹은 데이터 클래스와 같은 것을 PriorityQueue의 자료형으로 선언한다면 어떻게 될까?

 

따로 우선순위를 정해주지 않으면 에러가 나게 된다.

Pair를 자료형으로 했을 때, 첫번째 값으로 정렬을 할지, 두번째 값으로 정렬을 할지 모르기 때문에 에러가 나는 것이다.

이를 해결하기 위해서는 우선순위 기준을 정해줘야 한다.

import java.util.*

fun main() {
    val pq = PriorityQueue<Pair<Int, Int>>(Comparator { a, b -> a.first - b.first})
    pq.add(Pair(1, 1))
    pq.add(Pair(2, 2))
    pq.add(Pair(3, 3))

    println(pq.poll())
}

(1, 1)

위처럼 Comparator를 사용하여 첫 번째 값을 오름차순으로 정렬할 것이라면 first의 값을 통해 정렬해주고, 두 번째 값을 이용할 것이면 second를 사용하면 된다.

 

이때 Comparator는 생략하여도 된다.

import java.util.*

fun main() {
    val pq = PriorityQueue<Pair<Int, Int>>( { a, b -> a.first - b.first} )
    pq.add(Pair(1, 1))
    pq.add(Pair(2, 2))
    pq.add(Pair(3, 3))

    println(pq.poll())
}

(1, 1)

 

Triple일 때도 같은 방법으로 하면 된다. Pair와 다른 점으로는 Pair는 first와 second라면, Triple은 first, second, third라는 것이다.

 

이번엔 데이터 클래스를 통해 우선순위큐의 우선순위를 정해줄 것이다.

우선순위큐 또한 정렬 기준이 없다면 위와 같은 에러가 난다.

따라서 Comparable을 통해 기준을 정해줘야 한다.

import java.util.*

data class Test(val a: Int, val b: Int, val c: Int): Comparable<Test> {
    override fun compareTo(other: Test): Int = c - other.c
}

fun main() {
    val pq = PriorityQueue<Test>()
    pq.add(Test(1, 1, 1))
    pq.add(Test(2, 2, 2))
    pq.add(Test(3, 3, 3))

    println(pq.poll())
}

Test(a=1, b=1, c=1)