Find Median from Data Stream

HardHeap & Priority QueueHeapDesignTwo Pointers

Problem

Design a data structure that receives a stream of integers one at a time and can report the median of all numbers seen so far at any point. If an even number of values have been seen, the median is the average of the two middle values once sorted.

Example. After adding 1 and 2, the median is 1.5, the average of the two middles. After then adding 3, the median becomes 2, the single middle value of [1, 2, 3].

Key idea

Re-sorting all values seen so far every time the median is requested wastes the ordering already established by previous calls. The way to make this fast is to avoid maintaining a full order at all: only the boundary around the middle actually needs to stay correct.

Split the numbers seen so far into two halves: a max-heap holding the smaller half, a min-heap holding the larger half, kept balanced so their sizes differ by at most one. The max-heap's top is always the largest of the small half, and the min-heap's top the smallest of the large half. Together, these are the one or two middle values. A new number is inserted into the appropriate heap based on comparison with the max-heap's top, then the heaps are rebalanced by moving one top element across if one side has grown too large. The median reads directly off the top of one or both heaps.

Solution

// Minimal binary heap; the comparator (hasPriority) decides min-heap or max-heap order.
class Heap<T> {
  private readonly items: T[] = [];

  constructor(private readonly hasPriority: (a: T, b: T) => boolean) {}

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

  peek(): T {
    return this.items[0];
  }

  push(item: T): void {
    this.items.push(item);
    let i = this.items.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.hasPriority(this.items[parent], this.items[i])) {
        break;
      }
      [this.items[parent], this.items[i]] = [this.items[i], this.items[parent]];
      i = parent;
    }
  }

  pop(): T {
    const top = this.items[0];
    const last = this.items.pop()!;
    if (this.items.length > 0) {
      this.items[0] = last;
      this.bubbleDown(0);
    }
    return top;
  }

  private bubbleDown(start: number): void {
    let i = start;
    while (true) {
      const left = 2 * i + 1;
      const right = 2 * i + 2;
      let next = i;
      if (left < this.items.length && this.hasPriority(this.items[left], this.items[next])) {
        next = left;
      }
      if (right < this.items.length && this.hasPriority(this.items[right], this.items[next])) {
        next = right;
      }
      if (next === i) {
        break;
      }
      [this.items[i], this.items[next]] = [this.items[next], this.items[i]];
      i = next;
    }
  }
}

class MedianFinder {
  private readonly lower = new Heap<number>((a, b) => a > b); // max-heap: largest of the smaller half on top
  private readonly upper = new Heap<number>((a, b) => a < b); // min-heap: smallest of the larger half on top

  addNum(num: number): void {
    if (this.lower.size === 0 || num <= this.lower.peek()) {
      this.lower.push(num);
    } else {
      this.upper.push(num);
    }

    if (this.lower.size > this.upper.size + 1) {
      this.upper.push(this.lower.pop()); // rebalance: lower grew too big, shift its top over
    } else if (this.upper.size > this.lower.size) {
      this.lower.push(this.upper.pop()); // rebalance: keep lower at least as large as upper
    }
  }

  findMedian(): number {
    if (this.lower.size > this.upper.size) {
      return this.lower.peek(); // odd total: lower holds the extra middle element
    }
    return (this.lower.peek() + this.upper.peek()) / 2; // even total: average the two middle values
  }
}

Complexity

  • Time: O(log n) per insertion for the heap operations; O(1) to read the current median.
  • Space: O(n). Every inserted value is stored in one of the two heaps.

Watch out for

  • Keep the two heaps balanced after every insertion, not just occasionally, or the "top of each heap" invariant breaks and the median reads wrong.
  • When the heap sizes are equal, the median is the average of both tops; when unequal, it is the top of whichever heap holds one more element.

Pattern

This is the two-heap "balance point" pattern for a running median or percentile: split data into a lower and upper partition, each ordered inward, and keep them balanced so the boundary elements stay available in constant time. The same idea generalizes to sliding-window medians and other running-order-statistic problems.

Related questions