Binary heap implementation in kotlin

  • click + to see full implementation
import java.util.LinkedList

 //sampleStart
fun main() {

    fun addItemsToHeap(heap: Heap<Int>) {
        heap.apply {
            insert(10)
            insert(20)
            insert(40)
            insert(3)
            insert(100)
        }
    }

    val maxHeap = Heap<Int>(Comparator { o1: Int, o2: Int ->
        o1 - o2 // normal order, if o1 is smaller, return -ive
    })

    addItemsToHeap(maxHeap)
    println("maxHeap: $maxHeap")
    println("delete from maxHeap: ${maxHeap.delete()}")
    println("after deletion: $maxHeap")

    val minHeap = Heap<Int>(Comparator { o1: Int, o2: Int ->
        o2 - o1 // reverse order, if o2 is smaller, return -ive
    })

    println()
    addItemsToHeap(minHeap)
    println("minHeap: $minHeap")
    println("delete from minHeap: ${minHeap.delete()}")
    println("after deletion: $minHeap")
}
//sampleEnd



/**
 * Heap data structure, based on the comparator it can be min or max heap
 */
class Heap<T>(private val comparator: Comparator<T>) {
    private val tree = LinkedList<T>()

    fun insert(value: T) {
        // add new value to the end of tree list
        tree.add(value)

        // current index of the element from where we start traversing backwards
        // element was inserted at the end so this is last element index
        var currentIndex = tree.size.dec()

        // parent of current index node
        var parentIndex = getParentIndex(currentIndex)

        // loop until we reach the root node, parent of root is -1
        while (parentIndex != -1) {
            // parent value
            val parent = tree[parentIndex]
            // current value
            val current = tree[currentIndex]

            val comparison = comparator.compare(parent, current)

            // if parent is less than current value, swap!
            if (comparison < 0) {
                swap(parentIndex, currentIndex)
            } else {
                // otherwise, no swaps are required. exit the loop
                break
            }

            // set currentIndex to parent index
            currentIndex = parentIndex
            // set parent index of current index
            parentIndex = getParentIndex(currentIndex)
        }
    }

    /**
     * Delete the root node (max or min element in the heap) and return it.
     */
    fun delete(): T? {
        if (tree.isEmpty()) {
            return null // nothing to delete
        }
        val root = tree[getRootIndex()]

        val lastElement = tree.removeLast()

        if (tree.isEmpty()) {
            return root
        }

        tree[getRootIndex()] = lastElement

        // pull greater value to top (this can be minimum value based on comparator)

        // starting node
        var startParentIndex = getRootIndex()

        // left child index of starting node
        var leftChildIndex = getLeftIndex(startParentIndex)

        // right child index of starting node
        var rightChildIndex = getRightIndex(startParentIndex)

        while (true) {

            val leftChild = tree.getOrNull(leftChildIndex)
            val rightChild = tree.getOrNull(rightChildIndex)
            var newParentIndex = -1


            // if there are no more nodes to explore, quit!
            if (leftChild == null && rightChild == null) {
                break
            }

            // if one of child nodes is null, then it will be right node b/c
            // heap is represented by complete binary tree. and it will never happen
            // in a complete binary tree that left node is null while there is a right node.
            if (rightChild == null) {
                if (comparator.compare(tree[startParentIndex], leftChild) < 0) {
                    // set the parent index for next iteration and swap values
                    newParentIndex = leftChildIndex
                    swap(startParentIndex, leftChildIndex)
                }
            } else {
                // both nodes are present, compare them and pick the index for greater value node
                val greaterChildNodeIndex = if (comparator.compare(leftChild, rightChild) < 0) {
                    rightChildIndex
                } else {
                    leftChildIndex
                }

                // compare greater value node with current start node and swap values if
                // child node is greater than parent
                if (comparator.compare(tree[startParentIndex], tree[greaterChildNodeIndex]) < 0) {
                    // set the parent index for next iteration and swap values
                    newParentIndex = greaterChildNodeIndex
                    swap(startParentIndex, greaterChildNodeIndex)
                }
            }

            if (newParentIndex != -1) {
                // new parent index was set it means there was a swap
                // continue with child nodes
                startParentIndex = newParentIndex
                leftChildIndex = getLeftIndex(startParentIndex)
                rightChildIndex = getRightIndex(startParentIndex)
            } else {
                // there was no parent index set which means no swap happened.
                // break the loop
                break
            }
        }
        return root
    }

    /**
     * Return the top most element of heap - this is the element that will be deleted
     * if delete is invoked
     */
    fun peek(): T? = tree.getOrNull(getRootIndex())

    /**
     * @param node tree node
     * @return parent index of the given node
     */
    private fun getParentIndex(node: Int): Int = node.minus(1).floorDiv(2)

    /**
     * @param node tree node
     * @return left child index of the given node
     */
    private fun getLeftIndex(node: Int): Int = node.times(2).plus(1)

    /**
     * @param node tree node
     * @return right child index of the given node
     */
    private fun getRightIndex(node: Int): Int = node.times(2).plus(2)

    /**
     * Root will always be at index 1
     */
    private fun getRootIndex(): Int = 0

    private fun swap(indexA: Int, indexB: Int) {
        val tmp = tree[indexA]
        tree[indexA] = tree[indexB]
        tree[indexB] = tmp
    }

    override fun toString(): String {
        return tree.toString()
    }
}

top