Find Median from Data Stream

Examples

Input: addNum(1); addNum(2); findMedian()
Output: 1.5

The numbers so far are [1,2]. With an even count, the median is the average of both middle values: (1+2)/2 = 1.5.

Input: addNum(3); findMedian()
Output: 2

The numbers so far are [1,2,3]. With an odd count, the median is the single middle value, which is 2.

Constraints

• -10^5 <= num <= 10^5

• At most 5 * 10^4 calls total to addNum and findMedian

Approach

Keeping every number in a fully sorted array works, but inserting into the middle of a sorted array requires shifting elements - and you never actually need the whole array sorted, just always know what's near the middle.

The trick is to split the stream into two halves: a max-heap holding the smaller half of the numbers (so its top is the largest of the small half - the number just below the middle), and a min-heap holding the larger half (so its top is the smallest of the large half - the number just above the middle). Keeping the two heaps balanced in size (equal, or the max-heap one larger) means the median is always readable directly from their tops, with no need to look at anything else.

Every new number goes into whichever half it belongs to, and then the heaps are rebalanced by size if needed.

Solutions

Brute Force - Sorted Array with Insertion

Keep all numbers in a sorted array. Each addNum finds the correct insertion point and splices the number in; findMedian just reads the middle position(s) directly.

class MedianFinder {
  constructor() {
    this.sorted = [];
  }

  addNum(num) {
    let lo = 0;
    let hi = this.sorted.length;
    while (lo < hi) {
      const mid = (lo + hi) >> 1;
      if (this.sorted[mid] < num) lo = mid + 1;
      else hi = mid;
    }
    this.sorted.splice(lo, 0, num);
  }

  findMedian() {
    const n = this.sorted.length;
    const mid = n >> 1;
    if (n % 2 === 1) return this.sorted[mid];
    return (this.sorted[mid - 1] + this.sorted[mid]) / 2;
  }
}

Time: O(n) per addNum (binary search is O(log n), but splice's shift is O(n)); O(1) per findMedian · Space: O(n)

Optimal - Two Heaps

Split the stream across a max-heap (the smaller half) and a min-heap (the larger half), always kept balanced in size. The median is then read straight off their tops.

class Heap {
  constructor(compare) {
    this.data = [];
    this.compare = compare; // negative means a belongs above b
  }

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

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

  push(val) {
    this.data.push(val);
    let i = this.data.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.compare(this.data[i], this.data[parent]) < 0) {
        [this.data[i], this.data[parent]] = [this.data[parent], this.data[i]];
        i = parent;
      } else break;
    }
  }

  pop() {
    const top = this.data[0];
    const last = this.data.pop();
    if (this.data.length > 0) {
      this.data[0] = last;
      let i = 0;
      const n = this.data.length;
      while (true) {
        let best = i;
        const left = 2 * i + 1;
        const right = 2 * i + 2;
        if (left < n && this.compare(this.data[left], this.data[best]) < 0) best = left;
        if (right < n && this.compare(this.data[right], this.data[best]) < 0) best = right;
        if (best === i) break;
        [this.data[i], this.data[best]] = [this.data[best], this.data[i]];
        i = best;
      }
    }
    return top;
  }
}

class MedianFinder {
  constructor() {
    this.lowerHalf = new Heap((a, b) => b - a); // max-heap: smaller half of the numbers
    this.upperHalf = new Heap((a, b) => a - b); // min-heap: larger half of the numbers
  }

  addNum(num) {
    if (this.lowerHalf.size() === 0 || num <= this.lowerHalf.peek()) {
      this.lowerHalf.push(num);
    } else {
      this.upperHalf.push(num);
    }

    // Rebalance so lowerHalf has either the same count as upperHalf, or exactly one more.
    if (this.lowerHalf.size() > this.upperHalf.size() + 1) {
      this.upperHalf.push(this.lowerHalf.pop());
    } else if (this.upperHalf.size() > this.lowerHalf.size()) {
      this.lowerHalf.push(this.upperHalf.pop());
    }
  }

  findMedian() {
    if (this.lowerHalf.size() === this.upperHalf.size()) {
      return (this.lowerHalf.peek() + this.upperHalf.peek()) / 2;
    }
    return this.lowerHalf.peek();
  }
}

Time: O(log n) per addNum; O(1) per findMedian · Space: O(n) total across both heaps