Median of Data Stream Using Two Heaps in DSA Typescript

Mental Model

We keep two groups of numbers: smaller half and larger half, so the middle number is easy to find.

Analogy: Imagine two boxes: one holds smaller numbers, the other holds bigger numbers. The middle number is at the edge between these boxes.

MaxHeap (smaller half): 5 -> 3 -> 1
MinHeap (larger half): 6 -> 8 -> 9
Median is between top of MaxHeap and MinHeap.

Dry Run Walkthrough

Input: Stream: 5, 2, 8, 3

Goal: Find the median after each number is added

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

MaxHeap: [5]
MinHeap: []
Median: 5

Why: Start by putting first number in smaller half

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

MaxHeap: [2]
MinHeap: [5]
Median: (2 + 5) / 2 = 3.5

Why: Keep heaps balanced so median is average of tops

Step 3: Add 8 to min heap (larger half), then balance by moving smallest from min heap to max heap

MaxHeap: [5, 2]
MinHeap: [8]
Median: 5

Why: Max heap can have one extra element for odd count

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

MaxHeap: [3, 2]
MinHeap: [5, 8]
Median: (3 + 5) / 2 = 4

Why: Balance heaps to keep size difference at most one

Result:

MaxHeap: [3, 2]
MinHeap: [5, 8]
Median after last insertion: 4

Annotated Code

DSA Typescript

class Heap {
  data: number[] = [];
  comparator: (a: number, b: number) => boolean;

  constructor(comparator: (a: number, b: number) => boolean) {
    this.comparator = comparator;
  }

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

  peek(): number | null {
    return this.data.length === 0 ? null : this.data[0];
  }

  push(val: number): void {
    this.data.push(val);
    this.bubbleUp();
  }

  pop(): number | null {
    if (this.data.length === 0) return null;
    const top = this.data[0];
    const end = this.data.pop()!;
    if (this.data.length > 0) {
      this.data[0] = end;
      this.bubbleDown();
    }
    return top;
  }

  bubbleUp(): void {
    let index = this.data.length - 1;
    const element = this.data[index];
    while (index > 0) {
      const parentIndex = Math.floor((index - 1) / 2);
      const parent = this.data[parentIndex];
      if (this.comparator(element, parent)) {
        this.data[index] = parent;
        index = parentIndex;
      } else {
        break;
      }
    }
    this.data[index] = element;
  }

  bubbleDown(): void {
    let index = 0;
    const length = this.data.length;
    const element = this.data[0];
    while (true) {
      let leftChildIndex = 2 * index + 1;
      let rightChildIndex = 2 * index + 2;
      let swapIndex = -1;

      if (leftChildIndex < length) {
        if (this.comparator(this.data[leftChildIndex], element)) {
          swapIndex = leftChildIndex;
        }
      }

      if (rightChildIndex < length) {
        if (
          this.comparator(this.data[rightChildIndex],
          swapIndex === -1 ? element : this.data[leftChildIndex])
        ) {
          swapIndex = rightChildIndex;
        }
      }

      if (swapIndex === -1) break;

      this.data[index] = this.data[swapIndex];
      index = swapIndex;
    }
    this.data[index] = element;
  }
}

class MedianFinder {
  maxHeap: Heap; // smaller half
  minHeap: Heap; // larger half

  constructor() {
    this.maxHeap = new Heap((a, b) => a > b); // max heap
    this.minHeap = new Heap((a, b) => a < b); // min heap
  }

  addNum(num: number): void {
    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()!); // move largest from smaller half to larger half
    } else if (this.minHeap.size() > this.maxHeap.size()) {
      this.maxHeap.push(this.minHeap.pop()!); // move smallest from larger half to smaller half
    }
  }

  findMedian(): number | null {
    if (this.maxHeap.size() === 0 && this.minHeap.size() === 0) {
      return null; // no elements
    }
    if (this.maxHeap.size() > this.minHeap.size()) {
      return this.maxHeap.peek()!; // odd count, max heap top is median
    } else {
      return (this.maxHeap.peek()! + this.minHeap.peek()!) / 2; // even count, average of tops
    }
  }
}

// Driver code
const mf = new MedianFinder();
const stream = [5, 2, 8, 3];
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 keep smaller half and larger half

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

balance heaps if smaller half has more than one extra element

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

balance heaps if larger half has more elements

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

if odd count, median is top of smaller half

else { return (this.maxHeap.peek()! + this.minHeap.peek()!) / 2; }

if even count, median is average of tops of both heaps

Output

5

3.5

5

4

Complexity Analysis

Time: O(log n) per insertion because each addNum operation pushes and pops from heaps which take O(log n)

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

vs Alternative: Compared to sorting the entire stream after each insertion (O(n log n)), using two heaps is much faster for continuous median queries.

Edge Cases

empty stream

findMedian returns null or error if called before any insertion

DSA Typescript

if (this.maxHeap.size() === 0 && this.minHeap.size() === 0) { return null; }

all numbers equal

heaps balance normally, median is that number

DSA Typescript

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

stream with negative and positive numbers

heaps correctly separate smaller and larger halves, median computed correctly

DSA Typescript

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 access.

Common Mistakes

Mistake: Not balancing the heaps after insertion, causing incorrect median calculation
Fix: Always check and balance the heaps sizes after each insertion to keep size difference at most one

Mistake: Adding all numbers to one heap only
Fix: Use condition to decide which heap to add the number to keep halves separated

Summary

Maintains two heaps to quickly find the median of a data stream.

Use when you need fast median queries on a continuously growing list.

The median is always at the top of one or both heaps, so keep them balanced.

Practice

1. What is the main purpose of using two heaps (a max heap and a min heap) in the Median of Data Stream problem?

easy

A. To sort all numbers in ascending order

B. To store all numbers in a single data structure

C. To efficiently keep track of the lower half and upper half of the numbers separately

D. To remove duplicates from the data stream

Solution

1. Step 1: Understand the role of two heaps

   The max heap stores the smaller half of the numbers, and the min heap stores the larger half.

2. Step 2: Explain how this helps find the median

   By balancing these heaps, the median is either the top of one heap or the average of tops of both heaps.

3. Final Answer:

   To efficiently keep track of the lower half and upper half of the numbers separately -> Option C

4. Quick Check:

   Two heaps split data halves = efficient median [OK]

Hint: Two heaps split data halves for quick median [OK]

Common Mistakes:

• Thinking heaps sort all numbers fully
• Using one heap only
• Confusing heaps with arrays

2. Which TypeScript syntax correctly declares a max heap using a priority queue for numbers?

easy

A. const maxHeap = new PriorityQueue<number>((a, b) => b - a);

B. const maxHeap = new PriorityQueue<number>((a, b) => a - b);

C. const maxHeap = new PriorityQueue<number>();

D. const maxHeap = new PriorityQueue<number>((a, b) => a + b);

Solution

1. Step 1: Understand max heap comparator

   Max heap comparator returns positive if first is less than second, so b - a creates max heap.

2. Step 2: Check each option

   const maxHeap = new PriorityQueue<number>((a, b) => b - a); uses (b - a) which is correct for max heap; others do not create max heap order.

3. Final Answer:

   const maxHeap = new PriorityQueue<number>((a, b) => b - a); -> Option A

4. Quick Check:

   Comparator b - a = max heap [OK]

Hint: Max heap comparator subtracts first from second [OK]

Common Mistakes:

• Using a - b for max heap (creates min heap)
• Not providing comparator
• Using addition in comparator

3. Given the following sequence of numbers inserted into the data stream: [5, 15, 1, 3], what is the median after all insertions using two heaps?

medium

A. 3

B. 4

C. 5

D. 8

Solution

1. Step 1: Insert numbers and balance heaps

   Insert 5 -> maxHeap=[5] (top=5), minHeap=[]; Insert 15 -> minHeap=[15] (top=15), sizes 1:1; Insert 1 -> maxHeap=[5,1] (top=5), sizes 2:1; Insert 3 -> maxHeap=[5,3,1] (top=5), sizes 3:1 -> balance (max.size > min.size+1), move max.top=5 to minHeap -> maxHeap=[3,1] (top=3), minHeap=[5,15] (top=5), sizes 2:2.

2. Step 2: Calculate median from heaps

   Sizes equal, median = (maxHeap.top + minHeap.top)/2 = (3 + 5)/2 = 4.

3. Final Answer:

   4 -> Option B

4. Quick Check:

   Median of [5,15,1,3] = 4 [OK]

Hint: Median is average of two heap tops if sizes equal [OK]

Common Mistakes:

• Taking median as top of one heap only
• Not balancing heaps after insertion
• Wrong average calculation

4. Consider this TypeScript snippet for adding a number to the data stream median finder:

if (num <= maxHeap.peek()) {
  maxHeap.enqueue(num);
} else {
  minHeap.enqueue(num);
}

// Balance heaps
if (maxHeap.size() > minHeap.size() + 1) {
  minHeap.enqueue(maxHeap.dequeue());
} else if (minHeap.size() > maxHeap.size()) {
  maxHeap.enqueue(minHeap.dequeue());
}

What is the bug in this code?

medium

A. It compares num with minHeap.peek() instead of maxHeap.peek()

B. It balances heaps incorrectly by moving elements the wrong way

C. It uses enqueue instead of push for heaps

D. It does not handle empty maxHeap before peek() call

Solution

1. Step 1: Analyze peek() usage on maxHeap

   If maxHeap is empty, calling peek() causes error because no element exists yet.

2. Step 2: Identify missing empty check

   Code should check if maxHeap is empty before peek() or handle initial insertion separately.

3. Final Answer:

   It does not handle empty maxHeap before peek() call -> Option D

4. Quick Check:

   Empty heap peek() causes runtime error [OK]

Hint: Check heap empty before peek() to avoid errors [OK]

Common Mistakes:

• Assuming heaps always have elements
• Mixing enqueue and push methods
• Balancing heaps incorrectly

5. You want to extend the Median of Data Stream solution to support removing numbers from the stream efficiently. Which data structure modification is best to support both insertion and removal while maintaining median calculation?

hard

A. Use two heaps with a hash map to mark removed elements lazily

B. Use two balanced binary search trees (like AVL or Red-Black trees) instead of heaps

C. Use a single sorted array and binary search for insertion and removal

D. Use two queues to store lower and upper halves

Solution

1. Step 1: Understand removal challenge in heaps

   Heaps do not support efficient arbitrary removal, so direct removal is costly.

2. Step 2: Explain lazy removal with hash map

   Mark elements as removed in a hash map and skip them during heap operations to simulate removal efficiently.

3. Step 3: Compare other options

   Balanced trees support removal but are more complex; arrays are inefficient; queues do not maintain order.

4. Final Answer:

   Use two heaps with a hash map to mark removed elements lazily -> Option A

5. Quick Check:

   Lazy removal with hash map + heaps = efficient median update [OK]

Hint: Lazy removal with hash map keeps heaps efficient [OK]

Common Mistakes:

• Trying direct removal from heaps
• Using arrays causing slow operations
• Ignoring complexity of balanced trees