Find Median From Data Stream

Keep a running sorted list. Every addNum does a binary search for the insertion point (O(log n)) but then has to shift elements to make room (O(n)), so insertion is linear overall. findMedian just reads the middle index (or averages the two middle indices).

import bisect

class MedianFinder:
    def __init__(self):
        self.data = []

    def addNum(self, num: int) -> None:
        bisect.insort(self.data, num)

    def findMedian(self) -> float:
        n = len(self.data)
        mid = n // 2
        if n % 2 == 1:
            return float(self.data[mid])
        return (self.data[mid - 1] + self.data[mid]) / 2

This is simple and correct, but with up to 5 * 10^4 calls, the O(n) shifting cost per insertion adds up to O(n²) overall in the worst case.

Maintain two heaps that together hold every number seen:

• small: a max-heap (negated, since Python’s heapq is min-heap only) holding the smaller half of the numbers.
• large: a min-heap holding the larger half.

Kept balanced so len(small) is always either equal to len(large) or exactly one more. That invariant means the median is either small’s root alone (odd total count) or the average of both roots (even total count) — no scan required.

To insert a number while preserving both the balance invariant and the ordering invariant (every element in small is less than or equal to every element in large) in one motion: always push into small first, then immediately move small’s largest into large — this guarantees correct ordering even if the new number actually belonged in large. Then, if that left large bigger than small, move large’s smallest back.

import heapq

class MedianFinder:
    def __init__(self):
        self.small = []  # max-heap (store negatives), holds the smaller half
        self.large = []  # min-heap, holds the larger half

    def addNum(self, num: int) -> None:
        heapq.heappush(self.small, -num)
        # move small's max into large, guaranteeing small's max <= large's min
        heapq.heappush(self.large, -heapq.heappop(self.small))
        if len(self.large) > len(self.small):
            heapq.heappush(self.small, -heapq.heappop(self.large))

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

Tracing the example — addNum(1), addNum(2), findMedian(), addNum(3), findMedian() — watch the new number always enter small first, a value cross the seam whenever an invariant is threatened, and each median come straight off the two facing tops:

Both invariants survive every call, and each median was read directly off the heap tops — no scan, no comparisons beyond the top values.

Correctness: the push-then-shuffle sequence in addNum guarantees two invariants after every call: (1) every value in small is <= every value in large, and (2) the two heaps differ in size by at most 1, with small never smaller than large. Together these mean small’s root is always the true lower-median candidate and large’s root the upper-median candidate — reading them off directly gives the median without ever sorting.

Complexity: addNum does a constant number of heap pushes/pops, each O(log n). findMedian just reads one or two roots, O(1). Space is O(n) to store every number across both heaps.

Follow-ups: if every value is known to be in [0, 100], a counting array of size 101 tracking frequencies (plus a running total count) lets you walk it to find the middle index in O(100) = O(1) per call instead of O(log n). If 99% of values fall in [0, 100] but a few outliers don’t, keep that counting array for the common range plus a small fallback structure (e.g. a sorted list or the two-heap approach) for the rare out-of-range values, merging the two only when a query needs to cross the boundary.