find-median-data-stream.ts

import { MinHeap } from "../lib/min-heap.js";

/**
 * 295. Find Median from Data Stream  (Hard)
 * Link: https://leetcode.com/problems/find-median-from-data-stream/
 *
 * Support addNum(num) to ingest a stream of integers and findMedian() to return
 * the median of all values seen so far.
 *
 * Example:
 *   addNum(1); addNum(2); findMedian() -> 1.5
 *   addNum(3);            findMedian() -> 2
 *
 * Approach:
 *   Two heaps split the data at the median: a max-heap `low` for the smaller
 *   half and a min-heap `high` for the larger half. Keep their sizes balanced
 *   (low holds the extra element when the count is odd). The median is low's top
 *   (odd count) or the average of both tops (even count). The max-heap is a
 *   MinHeap keyed by the negated value.
 *
 * Time:  addNum O(log n), findMedian O(1).
 * Space: O(n)
 */
export class MedianFinder {
  private readonly low = new MinHeap<number>((x) => -x); // max-heap of smaller half
  private readonly high = new MinHeap<number>((x) => x); // min-heap of larger half

  addNum(num: number): void {
    // push to low, then shift its max over to high to keep order
    this.low.push(num);
    this.high.push(this.low.pop()!);

    // rebalance so low has equal or one more element than high
    if (this.high.size > this.low.size) {
      this.low.push(this.high.pop()!);
    }
  }

  findMedian(): number {
    if (this.low.size > this.high.size) {
      return this.low.peek()!;
    }
    return (this.low.peek()! + this.high.peek()!) / 2;
  }
}