295. Find Median from Data Stream

Problem¶

Design MedianFinder to add stream values and return the median of all values seen so far.

addNum(1) -> addNum(2) -> findMedian() = 1.5
addNum(3) -> findMedian() = 2.0

class MedianFinder:
    def __init__(self):
        pass

    def addNum(self, num: int) -> None:
        pass

    def findMedian(self) -> float:
        pass

Key trick¶

Maintain two heaps:

A max-heap for the smaller half.
A min-heap for the larger half.
Keep the max-heap the same size as the min-heap or one element larger.

Trap¶

Python provides only a min-heap, so values in the max-heap must be negated.
Heap sizes alone are insufficient; every value in the smaller half must also be no greater than every value in the larger half.
Maintaining a sorted list makes insertion \(O(n)\) because elements must be shifted.

Why is it interesting?¶

It turns a dynamic order-statistic problem into heap balancing, reducing insertion from \(O(n)\) to \(O(\log n)\) while keeping median queries at \(O(1)\).

Python solution¶

import heapq


class MedianFinder:
    def __init__(self):
        # low is a max-heap represented by negated values.
        # low has either the same size as high or one extra value.
        self.low = []
        self.high = []

    def addNum(self, num: int) -> None:
        # Insert into low, then move its maximum into high.
        heapq.heappush(self.low, -num)
        heapq.heappush(self.high, -heapq.heappop(self.low))

        # For an odd count, low must contain the extra value.
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))

    def findMedian(self) -> float:
        if len(self.low) > len(self.high):
            return float(-self.low[0])

        return (-self.low[0] + self.high[0]) / 2

Time: addNum is \(O(\log n)\); findMedian is \(O(1)\).
Space: \(O(n)\).

Comment on my solution¶

It is correct when findMedian is called after at least one insertion.
Binary search finds the insertion index in \(O(\log n)\), but list insertion still costs \(O(n)\) due to shifting elements.
The check below never detects an empty list because self.ordered is initialized as a list, not None.

if self.ordered is None:

If empty-state validation were needed, it should check whether the list is empty.
Over \(n\) insertions, the worst-case time is \(O(n^2)\), compared with \(O(n \log n)\) for the two-heap solution.
The frequency-table idea is appropriate for the follow-up where values are restricted to the fixed range from \(0\) through \(100\).

# Idea:
# - If all integer numbers from the stream are in the range `[0, 100]`, how would you optimize your solution?
#   - maybe a bucket sort
#   - no better using a frequency table

import bisect

# WORKS but slow
class MedianFinder:

    def __init__(self):
        self.ordered = []


    def addNum(self, num: int) -> None:
        idx = bisect.bisect_left(self.ordered, num)
        self.ordered.insert(idx, num)


    def findMedian(self) -> float:
        if self.ordered is None:
            raise ValueError("No numbers in the list.")

        n = len(self.ordered)
        mid = n // 2
        if n % 2:
            return self.ordered[mid]
        else:
            return (self.ordered[mid - 1] + self.ordered[mid]) / 2