본문 바로가기
코딩테스트 준비(kotlin)/자료구조

[프로그래머스 , kotlin] 이중우선순위큐

by 1chanhue1 2024. 8. 1.

import java.util.Collections
import java.util.PriorityQueue


fun solution(operations: Array<String>): IntArray {
    val minPriorityQueue = PriorityQueue<Int>()
    val maxPriorityQueue = PriorityQueue<Int>(Collections.reverseOrder())
    val entryCount = mutableMapOf<Int, Int>() //각 숫자의 개수를 저장하는 맵

    for (i in operations.indices) {
        val value = operations[i].split(" ")
        val operation = value[0]
        val num = value[1].toInt()

        when (operation) {
            "I" -> {
                minPriorityQueue.add(num)
                maxPriorityQueue.add(num)
                entryCount[num] = entryCount.getOrDefault(num, 0) + 1
            }

            "D" -> {
                if (num == 1) {
                    removeTop(maxPriorityQueue, entryCount)
                } else if (num == -1) {
                    removeTop(minPriorityQueue, entryCount)
                }
            }
        }
    }

    val result = IntArray(2)
    if (entryCount.values.sum() == 0) { // 비어있으면 0,0 반환하기
        result[0] = 0
        result[1] = 0
    } else {
        result[0] = getValidTop(maxPriorityQueue, entryCount)
        result[1] = getValidTop(minPriorityQueue, entryCount)
    }

    return result
}

private fun removeTop(heap: PriorityQueue<Int>, entryCount: MutableMap<Int, Int>) {
    while (heap.isNotEmpty()) {
        val top = heap.poll()
        if (entryCount.getOrDefault(top, 0) > 0) {
            entryCount[top] = entryCount[top]!! - 1
            break
        }
    }
}

private fun getValidTop(heap: PriorityQueue<Int>, entryCount: MutableMap<Int, Int>): Int {
    while (heap.isNotEmpty()) {
        val top = heap.poll()
        if (entryCount.getOrDefault(top, 0) > 0) {
            return top
        }
    }
    return 0
}

// 테스트 코드
fun main() {
    println(
        solution(
            arrayOf(
                "I 16",
                "I -5643",
                "D -1",
                "D 1",
                "D 1",
                "I 123",
                "D -1"
            )
        ).joinToString(", ")
    )
    println(
        solution(
            arrayOf(
                "I -45",
                "I 653",
                "D 1",
                "I -642",
                "I 45",
                "I 97",
                "D 1",
                "D -1",
                "I 333"
            )
        ).joinToString(", ")
    )
}