Design a running-median container
add(num) and median() with two heaps, and the return type that trips candidates up.
- Heaps and priority queues in object design — Scheduling, ranking and timers with a heap behind a clean interface.
- Designing method signatures — Inputs, outputs, errors and names that make misuse hard.
- Encapsulation — Hiding data behind methods so an object's rules cannot be broken from outside.
01The design, in the six-step framework
1. Requirements and use cases
One reported interview opened with "save numbers and find the median" as a design question, and the notes warn about the return type. The algorithm is known; the marks are for the contract and the invariants.
- A caller adds a number to the container in O(log n).
- A caller asks for the median of everything added so far in O(1).
- With an even count the median is the average of the two middle values, so the return type is a float even for integer input.
- Asking for the median of an empty container fails clearly.
- Follow-up: numbers are known to lie in 0 to 100; or a window keeps only the last k.
Out of scope: arbitrary percentiles (mentioned below) and persistence. Non-functional constraints: reads far outnumber writes in the usual framing, so median must be O(1); memory O(n), since an exact median of an unbounded stream needs all the values in the middle; and no float rounding surprises in the even case.
2. Core entities and relationships
- MedianFinder: the public object; owns two heaps and enforces the balance rule.
- Low heap: a max-heap of the smaller half; its top is the largest of the small values.
- High heap: a min-heap of the larger half; its top is the smallest of the large values.
- BucketMedian: the bounded-range variant; an array of 101 counts.
- WindowMedian: the sliding variant; heaps plus lazy deletion.
The finder composes both heaps; a caller never sees either. The variants are separate classes with the same two methods, sharing an interface so a caller can swap them. Python only has a min-heap, so the low heap stores negated values, and that detail is hidden inside the class by encapsulation.
3. Interfaces and key methods
from typing import Protocol
class MedianContainer(Protocol):
def add(self, x: float) -> None:
"""Record one value."""
def median(self) -> float:
"""Exact median as a float; raises ValueError when empty."""
Invariants for the heap version: every value in low is at most every value in high; the sizes satisfy len(low) == len(high) or len(low) == len(high) + 1. It follows that the median is the top of low when the count is odd and the mean of both tops when it is even.
max(low) ≤ min(high), and 0 ≤ |low| − |high| ≤ 1
4. Design patterns and the choices behind them
Two heaps split at the middle
The force is that only the middle matters, and a fully sorted list does more work than needed. A sorted list gives O(1) median and O(n) insert. Quickselect gives O(1) insert and O(n) median. Two heaps keep exactly the two middle candidates on top, for O(log n) insert and O(1) median. The trick is the insertion order: push into low, move low's top to high, then move high's top back if high grew larger. Every value crosses the boundary, so the ordering invariant holds without comparing against the median first.
The return type is part of the design
The force is the even case: the median of 1 and 2 is 1.5, so an integer return is a bug. In Java, (a + b) / 2 on ints truncates and can overflow, so the sum needs doubles. Writing -> float in the signature is the method design answer to the warning in the reports, and it should be raised before any code. An empty container raises instead of returning 0, because 0 is a valid median.
Bounded range: count, do not compare
The force is the follow-up "all values are 0 to 100". With 101 counters, add is O(1), and median walks at most 101 buckets, also O(1) in n. Memory stops growing with the stream. This variant is the reason for the shared interface: the choice becomes a constructor argument at the call site.
5. Core flow and code
Flow for adding 5, 15, 1 and asking for the median after each:
- Add 5: pushed to low, moved to high, then back to low because high is larger. low {5}, high {}. Median 5.
- Add 15: pushed to low, 15 moves to high. low {5}, high {15}. Sizes equal, median 10.0.
- Add 1: pushed to low, 5 moves to high, then high is larger so 5 moves back. low {1, 5}, high {15}. Median 5.
import bisect
import heapq
from collections import deque
from typing import Deque, List
class MedianFinder:
def __init__(self) -> None:
self.low: List[float] = [] # max-heap by negation
self.high: List[float] = [] # min-heap
def add(self, x: float) -> None:
heapq.heappush(self.low, -x)
heapq.heappush(self.high, -heapq.heappop(self.low))
if len(self.high) > len(self.low):
heapq.heappush(self.low, -heapq.heappop(self.high))
def median(self) -> float:
if not self.low:
raise ValueError("median of empty container")
if len(self.low) > len(self.high):
return float(-self.low[0])
return (-self.low[0] + self.high[0]) / 2
def __len__(self) -> int:
return len(self.low) + len(self.high)
class BucketMedian:
"""Values are integers in [0, limit); add is O(1), median is O(limit)."""
def __init__(self, limit: int = 101) -> None:
self.counts = [0] * limit
self.total = 0
def add(self, x: float) -> None:
self.counts[int(x)] += 1
self.total += 1
def _kth(self, k: int) -> int:
seen = 0
for value, count in enumerate(self.counts):
seen += count
if seen > k:
return value
raise IndexError(k)
def median(self) -> float:
if not self.total:
raise ValueError("median of empty container")
a = self._kth((self.total - 1) // 2)
b = self._kth(self.total // 2)
return (a + b) / 2
class WindowMedian:
"""Median of the last k values. Sorted list: O(k) add from the memmove, O(1) median."""
def __init__(self, k: int) -> None:
if k < 1:
raise ValueError("k must be at least 1")
self.k = k
self.window: Deque[float] = deque()
self.sorted: List[float] = []
def add(self, x: float) -> None:
self.window.append(x)
bisect.insort(self.sorted, x)
if len(self.window) > self.k:
old = self.window.popleft()
del self.sorted[bisect.bisect_left(self.sorted, old)]
def median(self) -> float:
n = len(self.sorted)
if not n:
raise ValueError("median of empty container")
return (self.sorted[(n - 1) // 2] + self.sorted[n // 2]) / 2
if __name__ == "__main__":
m = MedianFinder()
m.add(5)
assert m.median() == 5.0
m.add(15)
assert m.median() == 10.0
m.add(1)
assert m.median() == 5.0
m.add(2)
assert m.median() == 3.5 and isinstance(m.median(), float)
import random
rng = random.Random(4)
b, f = BucketMedian(), MedianFinder()
seen: List[int] = []
for _ in range(500):
v = rng.randrange(101)
b.add(v); f.add(v); seen.append(v)
s = sorted(seen)
n = len(s)
expect = (s[(n - 1) // 2] + s[n // 2]) / 2
assert b.median() == expect == f.median()
w = WindowMedian(3)
for v in [1, 9, 2, 8]:
w.add(v)
assert w.median() == 8.0
6. Extensions, concurrency and what interviewers probe
- The even case. Float return, double arithmetic, overflow-safe sum. Raise it before coding.
- Sliding window. The sorted list is simple and is O(k) per add. Two heaps with lazy deletion give O(log k), at the cost of more code; see sliding windows.
- Other percentiles. Keep the same two heaps with a different balance ratio, for example low holding 90 percent for p90. At scale, exact answers stop being worth their memory; approximate sketches take over, as in approximate counting.
- Memory limits. An exact median of an unbounded stream needs O(n) memory; if values are bounded, the bucket version escapes this. Say which case you are in.
- Thread safety. One lock over add and median; the two heaps are a single unit.
- Common mistake. Rebalancing by size only, without routing each value through the other heap, which breaks the ordering invariant silently.
02Where this question was reported
- LeetCode Discuss · LeetCode 295 (Hard), Google-tagged, freq 50.5; also in Prepfully's recent Google list
- Prepfully · Prepfully: 'median of a number stream' among recently reported Google coding questions
- LeetCode Discuss · 'Initially put as a design question'; mind the return type of findMedian()