Kotlin max heap example

max-heap-example.kt

data class Worker(
    val quality: Int, 
    val wage: Int, 
    val ratio: Double = wage.toDouble() / quality.toDouble()
)

class MaxHeap(private var capacity: Int = 32) {
    
    private var array = arrayOfNulls<Worker>(capacity)
    var size = 0
        private set
    
    fun offer(item: Worker) {
        ensureCapacity()
        array[size] = item
        size++
        heapifyUp()
    }
    
    fun poll(): Worker? {
        if (isEmpty) return null
        val item = array[0]
        array[0] = array[size - 1]
        size--
        heapifyDown()
        return item
    }
    
    fun peek(): Worker? = array[0]
    
    val isEmpty get() = (size == 0)
    
    val desc get() = "${array.filter { it != null }}"
    
    private fun heapifyDown() {
        var idx = 0
        while (hasLeftChild(idx)) {
            var smallerChildIdx = leftChildIndex(idx)
            if (hasRightChild(idx)  && rightChild(idx)!!.quality > leftChild(idx)!!.quality) smallerChildIdx = rightChildIndex(idx)
            if (array[idx]!!.quality > array[smallerChildIdx]!!.quality) {
                break
            } else {
                swap(idx, smallerChildIdx)
            }
            idx = smallerChildIdx
            
        }
    }
    
    private fun heapifyUp() {
        var idx = size - 1
        while (hasParent(idx) && parent(idx)!!.quality < array[idx]!!.quality) {
            swap(parentIndex(idx), idx)
            idx = parentIndex(idx)
        }
    }
    
    private fun swap(one: Int, two: Int) {
        val temp = array[one]
        array[one] = array[two]
        array[two] = temp
    }
    
    private fun ensureCapacity() {
        if (size == capacity) {
            array = Array<Worker?>(capacity * 2) { null }.apply {
                array.forEachIndexed { i, w -> this[i] = w}
            }
            capacity *= 2
        }
    }
    
    private fun leftChildIndex(parentIdx: Int) = (2 * parentIdx + 1)
    private fun rightChildIndex(parentIdx: Int) = (2 * parentIdx + 2)
    private fun parentIndex(childIdx: Int) = ((childIdx - 1) / 2)
    private fun hasLeftChild(idx: Int) = leftChildIndex(idx) < size
    private fun hasRightChild(idx: Int) = rightChildIndex(idx) < size
    private fun hasParent(idx: Int) = parentIndex(idx) >= 0
    private fun leftChild(idx: Int) = array[leftChildIndex(idx)]
    private fun rightChild(idx: Int) = array[rightChildIndex(idx)]
    private fun parent(idx: Int) = array[parentIndex(idx)]
    
}