Find Median from Data Stream

Signal

"Running median over a stream" needs constant-time access to both the largest of the lower half and the smallest of the upper half at once — one heap can't give you two different ends cheaply, so you split the data across two heaps and keep them balanced.

Approach

Maintain a max-heap (small, negated) holding the lower half of the numbers and a min-heap (large) holding the upper half, kept so every value in small is ≤ every value in large, and their sizes differ by at most one. On each insert, push into small, then shuffle its max over to large to preserve the ordering invariant, then — if that left large bigger — shuffle one back so small never falls behind. The median is then either the top of whichever heap is larger, or the average of both tops when they're equal size.

Skeleton

small = max-heap (lower half)   # negate values
large = min-heap (upper half)

def add(num):
    push small, num
    push large, pop(small)          # rebalance ordering
    if len(large) > len(small):
        push small, pop(large)      # rebalance size

def median():
    if len(small) > len(large): return -small[0]
    return (-small[0] + large[0]) / 2

Solution

import heapq

class MedianFinder:
    def __init__(self):
        self.small: list[int] = []  # max-heap, values negated
        self.large: list[int] = []  # min-heap

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

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

Complexity

O(log n) time per add_num (two heap operations), O(1) for find_median. O(n) total space across both heaps.

Pitfalls

• Pushing straight into small or large without the cross-push rebalancing step breaks the "every small ≤ every large" invariant, which silently corrupts the median even though sizes stay balanced.
• Allowing the size difference to grow past 1 — always route through the push-then-rebalance dance on every insert, not just when it "looks unbalanced".
• Re-sorting the whole dataset on every find_median call also works but is O(n log n) per query — the two-heap approach amortizes insertion cost instead.