[kotlin] 코틀린으로 힙 알고리즘 구현하기

14 Dec 2023

kotlin

힙 알고리즘은 이진트리 구조를 사용하여 효율적으로 최솟값 또는 최댓값을 찾는 알고리즘입니다. 힙 알고리즘은 우선순위 큐 등 많은 응용분야에서 사용됩니다. 이번에는 코틀린으로 간단한 힙 알고리즘을 구현해보겠습니다.

힙 구현하기

class MinHeap {
    private val heapArray = ArrayList<Int>()

    fun add(value: Int) {
        heapArray.add(value)
        heapifyUp(heapArray.size - 1)
    }

    private fun heapifyUp(index: Int) {
        var currentIndex = index
        var parentIndex = (currentIndex - 1) / 2
        while (currentIndex > 0 && heapArray[currentIndex] < heapArray[parentIndex]) {
            heapArray.swap(currentIndex, parentIndex)
            currentIndex = parentIndex
            parentIndex = (currentIndex - 1) / 2
        }
    }

    fun removeMin(): Int? {
        if (heapArray.isEmpty()) return null
        if (heapArray.size == 1) return heapArray.removeAt(0)
        
        val minItem = heapArray[0]
        heapArray[0] = heapArray.removeAt(heapArray.size - 1)
        heapifyDown(0)
        return minItem
    }

    private fun heapifyDown(index: Int) {
        var currentIndex = index
        while (true) {
            val leftChildIndex = 2 * currentIndex + 1
            val rightChildIndex = 2 * currentIndex + 2
            var smallestIndex = currentIndex

            if (leftChildIndex < heapArray.size && heapArray[leftChildIndex] < heapArray[smallestIndex]) {
                smallestIndex = leftChildIndex
            }
            if (rightChildIndex < heapArray.size && heapArray[rightChildIndex] < heapArray[smallestIndex]) {
                smallestIndex = rightChildIndex
            }
            if (smallestIndex == currentIndex) return

            heapArray.swap(currentIndex, smallestIndex)
            currentIndex = smallestIndex
        }
    }

    private fun ArrayList<Int>.swap(index1: Int, index2: Int) {
        val temp = this[index1]
        this[index1] = this[index2]
        this[index2] = temp
    }
}

위 코드는 최소 힙을 구현한 예시입니다. MinHeap 클래스를 사용하여 원소를 추가하고 최솟값을 제거할 수 있습니다. add 메서드는 원소를 추가한 후에 heapifyUp 함수를 호출하여 적절한 위치로 이동시키고, removeMin 메서드는 힙의 루트를 제거한 후에 heapifyDown 함수를 호출하여 새로운 최솟값을 찾습니다.

테스트하기

fun main() {
    val minHeap = MinHeap()

    minHeap.add(10)
    minHeap.add(5)
    minHeap.add(8)
    minHeap.add(3)

    println(minHeap.removeMin())  // 출력: 3
    println(minHeap.removeMin())  // 출력: 5
}

위의 예제 코드는 MinHeap 클래스를 사용하여 최소 힙을 생성하고 테스트하는 예시입니다. add 메서드로 값을 추가한 후 removeMin 메서드를 이용하여 최솟값을 순서대로 출력합니다.

힙 알고리즘은 다양한 응용 분야에서 유용하게 활용되며, 위에서 소개한 예시 코드를 통해 코틀린으로 힙 알고리즘을 구현하는 법을 배워보았습니다.

참고 자료

코틀린 공식 문서