Trying out MaxHeaps

MaxHeap.kt

package com.aniketkadam.heaps

import org.jetbrains.annotations.TestOnly

class MaxHeap<T : Comparable<T>> {
    private val items: MutableList<T> = mutableListOf()

    fun insertAll(items: List<T>): Unit = items.forEach(::insert)

    fun insert(item: T): Unit {
        items.add(item)
        bubbleUp(items.size - 1)
    }

    private fun getParentIndex(index: Int): Int {
        // Left child = 2i + 1
        // Right child = 2i + 2
        // So the parent of 'i' is:
        //  2pi + 1 = ci
        // pi = Integer value of (c1 - 1)/2
        return (index - 1) / 2
    }

    private fun bubbleUp(index: Int) {
        // If it's the root item, it'll be compared to itself and will never be greater than itself
        val parentIndex = getParentIndex(index)
        val isGreaterThanParent = items[index] > items[parentIndex]

        if (isGreaterThanParent) {
            // Swap with the parent
            val temp = items[index]
            items[index] = items[parentIndex]
            items[parentIndex] = temp
            // Check the parent
            bubbleUp(parentIndex)
        }

    }

    private fun bubbleDown(index: Int = 0) {
        // Get the left and right child
        // Left child = 2i + 1
        // Right child = 2i + 2
        val leftIdx = 2 * index + 1
        val rightIdx = 2 * index + 2

        var compareIndex = -1
        if (rightIdx < items.size) {
            // Both left and right children exist
            // The compare index is the one which is greater
            if(items[leftIdx] > items[rightIdx]) {
                compareIndex = leftIdx
            } else {
                compareIndex = rightIdx
            }
        } else if (leftIdx < items.size) {
            // Only left child exists
            compareIndex = leftIdx
        } else {
            // There are no children so nothing to do
            return
        }

        // If we get here then compareIndex is not -1
        //  Also a swap may be necessary
        if(items[index] < items[compareIndex]) {
            val temp = items[index]
            items[index] = items[compareIndex]
            items[compareIndex] = temp
            bubbleDown(compareIndex)
        }
    }

    fun extract(): T {
        val max = items[0]
        items[0] = items.removeAt(items.size - 1)
        bubbleDown()
        return max
    }

    @TestOnly
    fun shape(): List<T> = items
}

MaxHeapTest.kt

package com.aniketkadam.heaps

import org.hamcrest.core.IsEqual.equalTo
import org.junit.Test

import org.junit.Assert.*
import org.junit.Before

class MaxHeapTest {

    lateinit var heap: MaxHeap<Int>

    @Before
    fun setup() {
        heap = MaxHeap()
    }

    @Test
    fun `insert works as expected`() {
        heap.insertAll(listOf(1, 2, 3, 4, 5))
        assertThat(heap.shape(), equalTo(listOf(5, 4, 2, 1, 3)))
    }

    @Test
    fun `extract works as expected`() {
        heap.insertAll(listOf(1, 2, 3, 4, 5))
        assertThat(heap.extract(), equalTo(5))
        assertThat(heap.shape(), equalTo(listOf(4, 3, 2, 1)))
    }
}