Median of Data Stream Using Two Heaps in DSA Javascript

Mental Model

We keep two heaps to split numbers into smaller half and larger half so we can quickly find the middle value.

Analogy: Imagine sorting people by height into two lines: one line for shorter people and one for taller people. The middle person is between these two lines.

MaxHeap (smaller half)       MinHeap (larger half)
      [5, 3, 1]                  [6, 8, 9]
       ↑ top is max              ↑ top is min

Dry Run Walkthrough

Input: Stream: 5, 3, 8, 9, 1, 6

Goal: Find the median after each new number arrives

Step 1: Add 5 to max heap (smaller half)

MaxHeap: [5] ↑top
MinHeap: []

Why: Start by putting first number in smaller half

Step 2: Add 3 to max heap, then balance by moving max from max heap to min heap

MaxHeap: [3] ↑top
MinHeap: [5] ↑top

Why: Keep smaller half max heap larger or equal size than min heap

Step 3: Add 8 to min heap, then balance by moving min from min heap to max heap

MaxHeap: [5, 3] ↑top
MinHeap: [8] ↑top

Why: Balance heaps so size difference is at most 1

Step 4: Add 9 to min heap, no balancing needed

MaxHeap: [5, 3] ↑top
MinHeap: [8, 9] ↑top

Why: Heaps sizes are now equal

Step 5: Add 1 to max heap, no balancing needed

MaxHeap: [5, 3, 1] ↑top
MinHeap: [8, 9] ↑top

Why: Max heap can be one larger than min heap

Step 6: Add 6 to min heap, no balancing needed

MaxHeap: [5, 3, 1] ↑top
MinHeap: [6, 8, 9] ↑top

Why: Heaps are balanced with sizes equal

Result:

MaxHeap: [5, 3, 1] ↑top
MinHeap: [6, 8, 9] ↑top
Median after each insertion: 5, 4, 5, 6.5, 5, 5.5

Annotated Code

DSA Javascript

class Heap {
  constructor(compare) {
    this.data = [];
    this.compare = compare;
  }

  size() {
    return this.data.length;
  }

  peek() {
    return this.data[0];
  }

  push(val) {
    this.data.push(val);
    this._heapifyUp();
  }

  pop() {
    if (this.size() === 0) return null;
    const top = this.data[0];
    const end = this.data.pop();
    if (this.size() > 0) {
      this.data[0] = end;
      this._heapifyDown();
    }
    return top;
  }

  _heapifyUp() {
    let idx = this.size() - 1;
    while (idx > 0) {
      let parent = Math.floor((idx - 1) / 2);
      if (this.compare(this.data[idx], this.data[parent])) {
        [this.data[idx], this.data[parent]] = [this.data[parent], this.data[idx]];
        idx = parent;
      } else {
        break;
      }
    }
  }

  _heapifyDown() {
    let idx = 0;
    const length = this.size();
    while (true) {
      let left = 2 * idx + 1;
      let right = 2 * idx + 2;
      let swap = idx;

      if (left < length && this.compare(this.data[left], this.data[swap])) {
        swap = left;
      }
      if (right < length && this.compare(this.data[right], this.data[swap])) {
        swap = right;
      }
      if (swap === idx) break;
      [this.data[idx], this.data[swap]] = [this.data[swap], this.data[idx]];
      idx = swap;
    }
  }
}

class MedianFinder {
  constructor() {
    this.maxHeap = new Heap((a, b) => a > b); // smaller half
    this.minHeap = new Heap((a, b) => a < b); // larger half
  }

  addNum(num) {
    if (this.maxHeap.size() === 0 || num <= this.maxHeap.peek()) {
      this.maxHeap.push(num); // add to smaller half
    } else {
      this.minHeap.push(num); // add to larger half
    }

    // Balance heaps sizes
    if (this.maxHeap.size() > this.minHeap.size() + 1) {
      this.minHeap.push(this.maxHeap.pop());
    } else if (this.minHeap.size() > this.maxHeap.size()) {
      this.maxHeap.push(this.minHeap.pop());
    }
  }

  findMedian() {
    if (this.maxHeap.size() > this.minHeap.size()) {
      return this.maxHeap.peek();
    } else {
      return (this.maxHeap.peek() + this.minHeap.peek()) / 2;
    }
  }
}

// Driver code
const mf = new MedianFinder();
const stream = [5, 3, 8, 9, 1, 6];
for (const num of stream) {
  mf.addNum(num);
  console.log(mf.findMedian());
}

if (this.maxHeap.size() === 0 || num <= this.maxHeap.peek()) {

decide which heap to add the new number to based on comparison with max of smaller half

this.maxHeap.push(num);

add number to max heap (smaller half)

this.minHeap.push(num);

add number to min heap (larger half)

if (this.maxHeap.size() > this.minHeap.size() + 1) {

check if max heap is too big and rebalance

this.minHeap.push(this.maxHeap.pop());

move top from max heap to min heap to balance sizes

else if (this.minHeap.size() > this.maxHeap.size()) {

check if min heap is bigger and rebalance

this.maxHeap.push(this.minHeap.pop());

move top from min heap to max heap to balance sizes

if (this.maxHeap.size() > this.minHeap.size()) {

if max heap bigger, median is top of max heap

return (this.maxHeap.peek() + this.minHeap.peek()) / 2;

if equal size, median is average of tops

Output

5

4

5

6.5

5

5.5

Complexity Analysis

Time: O(log n) per insertion because each addNum operation inserts into a heap and balances heaps which takes log n time

Space: O(n) because all numbers are stored in the two heaps

vs Alternative: Compared to sorting all numbers each time (O(n log n)), this approach is faster for streaming data as it keeps partial order in heaps

Edge Cases

empty stream

findMedian returns undefined or error if called before any number is added

DSA Javascript

if (this.maxHeap.size() === 0 || num <= this.maxHeap.peek()) {

all numbers equal

median is that number, heaps balance correctly

DSA Javascript

if (this.maxHeap.size() > this.minHeap.size() + 1) {

stream with one number

median is that single number

DSA Javascript

if (this.maxHeap.size() > this.minHeap.size()) {

When to Use This Pattern

When you need to find the median of a growing list of numbers efficiently, use two heaps to keep track of smaller and larger halves for quick median retrieval.

Common Mistakes

Mistake: Not balancing the heaps after insertion causing incorrect median calculation
Fix: Always rebalance heaps so their sizes differ by at most one after each insertion

Mistake: Adding all numbers to one heap only
Fix: Distribute numbers between max heap and min heap based on comparison with current median

Summary

Keeps track of median in a stream by using two heaps to split numbers into smaller and larger halves.

Use when you want to find median quickly after each new number in a stream.

The key is balancing two heaps so median is always at the top of one or average of tops.