Find Median from Data Stream

hard · heap, data-structure

Find Median from Data Stream

Design a data structure that supports adding numbers from a stream and finding the median at any time.

Why this is worth learning

Split the multiset into a lower half and an upper half. A max-heap exposes the lower half’s largest value; a min-heap exposes the upper half’s smallest value.

Before coding: one-minute experiment

After inserting 5,1,9,2, what are the middle two values and median? Why can two independently rounded huge operands lose a small exact mean?

Function signature

type MedianFinder struct { /* fields */ }

func Constructor() MedianFinder
func (m *MedianFinder) AddNum(num int)
func (m *MedianFinder) FindMedian() float64

Inputs / outputs

  • AddNum: inserts a number into the data structure.
  • FindMedian: returns the current median.

Example

AddNum(1), AddNum(2) -> FindMedian() = 1.5
AddNum(3) -> FindMedian() = 2

Constraints

  • 1 <= number of operations <= 100_000

Core idea (near-solution)

Sorting every prefix makes the middle values visible but repeatedly orders old data. Keeping only one current median also loses information: the next insertion may move the rank boundary to a value that was discarded. Preserve both sides of that boundary instead. Only the maximum of the lower side and minimum of the upper side are needed for a query, which motivates two oppositely ordered heaps.

Maintain two heaps: a max-heap for the lower half and a min-heap for the upper half. Balance them so their sizes differ by at most one.

Invariant: every low value is <= every high value, and low has exactly the same number of entries as high or one more.

Algorithm (step-by-step)

  1. AddNum(x):
    • If low is empty or x <= max(low), push into low; else push into high.
    • Rebalance if len(low) > len(high)+1 or len(high) > len(low).
  2. FindMedian():
    • If sizes equal, return the safely computed arithmetic mean of max(low) and min(high), rounded to float64. Never add two signed ints before checking their signs.
    • Otherwise return max(low).

Correctness sketch

Maintain every low value <= every high value, with low having either the same number of entries as high or exactly one more. Insert relative to low’s root, then move one boundary root if sizes violate this rule. Because the moved root is the extreme at the boundary, transferring it preserves ordering while restoring sizes. With odd total size, low’s root is the middle-ranked value. With even size, low’s maximum and high’s minimum are the two middle values, whose exact arithmetic mean is returned as float64. Independent instances own independent heaps, and querying the median does not consume either root.

Trace (hand walkthrough)

Trace Walkthrough
Example input:
AddNum(1), AddNum(2) -> FindMedian() = 1.5
AddNum(3) -> FindMedian() = 2

Walk:
- Add 1 -> low=[1], high=[], median=1
- Add 2 -> low=[1], high=[2], median=(1+2)/2=1.5
- Add 3 -> rebalance low=[2,1], high=[3], median=2

Pitfalls to avoid

  • Failing to rebalance after insertion.
  • Returning the wrong heap's top when sizes differ.
  • Using a min-heap for low by mistake.

Debugging lens

Test the median after every insertion, not just at the end. Include duplicates, decreasing streams, two independent instances, and signed extrema. FindMedian requires at least one inserted value; querying an empty structure is not a supported operation.

Complexity

  • Time per AddNum: O(log n).
  • Time per FindMedian: O(1).
  • Extra space: O(n).

Implementation notes (Go)

  • Use container/heap with a custom max-heap for low and min-heap for high.
  • Keep values as ints. For the even median, opposite signs permit safe signed addition before conversion; same-sign sums need unsigned magnitudes or another exact wider representation. Converting both operands first can lose cancellation (MinInt and MaxInt should average to -0.5).

Contract precision

Use Constructor before AddNum, and call FindMedian only after at least one insertion. The corrected mean helper avoids both signed overflow and cancellation loss across the int domain.

Reconstruction checkpoint

After inserting 5,1,9,2, what are the middle two values and median? Why can two independently rounded huge operands lose a small exact mean?

Model answer

The sorted values are [1,2,5,9], so the middle values are 2 and 5 and the median is 3.5. Near opposite int extrema, converting each operand first can round away their one-unit difference; summing safely in integer form before the final scale preserves a result such as -0.5.

Return to Heap / Priority Queue · Module quiz

Run tests to see results
No issues detected
    Join Discord