Find Median From Data Stream

Add numbers one at a time and return the current median.

Start with the concrete trace below. It shows the state the algorithm must carry as it runs.

Problem trace

Find Median From Data Stream: The lower max-heap owns the extra item, and its root plus the upper min-heap root are always the middle value or values.

Input and goalAdd numbers one at a time and return the current median.

Initialize two empty heapslower and upper both start empty. Python stores negated values in lower so its minimum stored key represents the largest lower-half value.

heap	stored values	logical role
lower	[]	max-heap via negatives
upper	[]	min-heap

stream[5,2,10,4]

Recognize it

Use two heaps when numbers arrive online and every median query must avoid sorting the entire history.

Keep true

Every lower value is at most every upper value, and lower has either the same size as upper or exactly one extra item. Therefore the middle values are always heap roots.

Reuse it

Maintain an ordered partition whose size difference is bounded, then answer rank queries from boundary roots. The same balancing idea supports streaming quantiles and running order statistics.

Pattern: Two heaps.

Simple idea: A max-heap stores the lower half and a min-heap stores the upper half. Keep the lower half the same size as the upper half or one item larger. The middle value or values are then at the heap roots.

Python has only a min-heap, so negative values create the max-heap.

import heapq

class MedianFinder:
   def __init__(self) -> None:
      self.lower: list[int] = []
      self.upper: list[int] = []

   def add_num(self, num: int) -> None:
      heapq.heappush(self.lower, -num)
      heapq.heappush(self.upper, -heapq.heappop(self.lower))
      if len(self.upper) > len(self.lower):
         heapq.heappush(self.lower, -heapq.heappop(self.upper))

   def find_median(self) -> float:
      if len(self.lower) > len(self.upper):
         return float(-self.lower[0])
      return (-self.lower[0] + self.upper[0]) / 2

Cost: Adding takes $O(\log n)$ time. Finding the median takes $O(1)$ time. Space is $O(n)$.