Elwin's Blog

an activist who likes to think

front cover

Implement a MinHeap in Javascript

Posted on

It's useful when looking for a subset with certain quality in a stream, like Kth largest Element or Kth smallest Element.

The keypoint is understand how to get the parent index when bubble up, and left child index and right child index when trickle down, below is the implementation with comments:

class minHeap {
    constructor(capacity) {
        this.size = capacity;
        this.minHeap = []
    }
    
    add(val) {
        if (this.minHeap.length < this.size) {
            this.minHeap.push(val)
            // it's possible that the new element can be smaller than 
            // some element in this heap, so call in bubbleUp function
            this.bubbleUp(this.minHeap.length - 1)
            // if the new element is greater than the smallest element 
            // in a MinHeap, it will replace one of the element in the MinHeap
        } else if (val > this.minHeap[0]) {
            this.minHeap[0] = val
            this.trickleDown(0)
        }
    }
    
    peak() {
        return this.minHeap[0]
    }
    
    bubbleUp(index) {
        const parentIndex = Math.floor((index - 1) / 2)
        
        let max = index
        // if this.minHeap[parentIndex] > this.minHeap[max], than the element 
        // been inserted should bubble up
        if (parentIndex >= 0 && this.minHeap[parentIndex] > this.minHeap[max]) {
            max = parentIndex
        }
        
        if (max !== index) {
            this.swap(max, index)
            this.bubbleUp(max)
        }
    }
    
    trickleDown(index) {
        let leftChildIndex = index * 2 + 1
        let rightChildIndex = index * 2 + 2
        
        let min = index
        // if element at min index is bigger than element at leftChildIndex 
        // or rightChild Index, it should goes down the heap
        if (leftChildIndex < this.minHeap.length && this.minHeap[leftChildIndex] < this.minHeap[min]) {
            min = leftChildIndex
        }
        
        if (rightChildIndex < this.minHeap.length && this.minHeap[rightChildIndex] < this.minHeap[min]) {
            min = rightChildIndex
        }
        
        if (min !== index) {
            this.swap(min, index)
            this.trickleDown(min)
        }
    }
    
    swap(a, b) {
        [this.minHeap[a], this.minHeap[b]] = [this.minHeap[b], this.minHeap[a]]
    }
}

When finding Kth largest Element in a stream, it can be used like this:

var KthLargest = function(k, nums) {
    this.heap = new minHeap(k)
    
    for (let num of nums) {
        this.heap.add(num)
    }
};

KthLargest.prototype.add = function(val) {
    this.heap.add(val)
    return this.heap.peak()
};

It's lots of code, but once I grasp the core of a MinHeap, it's not too hard to write it.