Find Median from Data Stream

This is the showpiece of the pattern, and it is a design question: you build a class, not a single function. One heap cannot find a middle value. Two heaps, facing each other, can.

The problem

Build a class with two methods. add_num(num) adds an integer from a stream. find_median() returns the median of all numbers added so far: the middle value if the count is odd, or the average of the two middle values if it is even.

• Add 5 → median 5.0. Add 15 → median 10.0 (the average of 5 and 15). Add 1 → median 5.0 (sorted 1, 5, 15). Add 3 → median 4.0 (sorted 1, 3, 5, 15, and (3 + 5) / 2 = 4).

Constraints: up to 5 × 10⁴ calls in total, values between -10⁵ and 10⁵, and find_median is only called after at least one add_num.

Clarifying questions

• How often is each method called? Assume both are frequent, so both must be fast.
• Return type for an even count? A float, the average of the two middle values.
• Is there a range on the values? Yes here, but don't depend on it unless asked — it is a follow-up.

Approach 1: a sorted list

Keep every number in a sorted list. The median is then one or two index reads.

1import bisect23class MedianFinderSorted:4    """Keep every value in a sorted list."""56    def __init__(self) -> None:7        self.values: list[int] = []89    def add_num(self, num: int) -> None:10        bisect.insort(self.values, num)          # O(log n) search, O(n) shift1112    def find_median(self) -> float:13        n = len(self.values)14        middle = n // 215        if n % 2:16            return float(self.values[middle])17        return (self.values[middle - 1] + self.values[middle]) / 2

find_median is O(1). But add_num is O(n): binary search finds the spot in O(log n), and then every larger number shifts one place right. Over 5 × 10⁴ inserts that is up to about 1.25 × 10⁹ element moves in the worst case. Python's shift is a fast memory copy, so this may pass, but it is the wrong complexity, and interviewers ask for O(log n) inserts. Sorting on every query instead is worse: O(n log n) per call.

The key insight

The median only depends on the boundary between the smaller half and the larger half of the data. You never need the order inside either half.

So split the data in two:

• lower, a max-heap holding the smaller half. Its root is the largest of the small numbers.
• upper, a min-heap holding the larger half. Its root is the smallest of the large numbers.

The two roots face each other across the middle. Keep two invariants after every insert:

1. Order: everything in lower is at most everything in upper.
2. Balance: the sizes are equal, or lower has exactly one more.

Then the median is lower's root when the count is odd, and the average of both roots when it is even. Both roots are O(1) to read.

Approach 2: two heaps

1import heapq23class MedianFinder:4    """Smaller half in a max-heap, larger half in a min-heap."""56    def __init__(self) -> None:7        self.lower: list[int] = []   # max-heap of the smaller half (negated values)8        self.upper: list[int] = []   # min-heap of the larger half910    def add_num(self, num: int) -> None:11        heapq.heappush(self.lower, -num)                       # 1. enter via lower12        heapq.heappush(self.upper, -heapq.heappop(self.lower)) # 2. lower's max moves up13        if len(self.upper) > len(self.lower):                  # 3. restore the sizes14            heapq.heappush(self.lower, -heapq.heappop(self.upper))1516    def find_median(self) -> float:17        if len(self.lower) > len(self.upper):18            return float(-self.lower[0])                       # odd count19        return (-self.lower[0] + self.upper[0]) / 2            # even count

The three steps are the heart of it:

1. Push the new number into lower.
2. Move lower's largest into upper. The new number has now been compared against the boundary, so it is on the correct side. This step enforces order.
3. If upper is now bigger than lower, move upper's smallest back. This step enforces balance, and moving the smallest of the upper half cannot break order.

Dry run on the stream 5, 15, 1, 3. Heaps are shown as sets of real values (lower stores them negated).

Check the last row: sorted, the four numbers are 1, 3, 5, 15, and the median is 4. Correct.

Complexity. add_num does at most three pushes and two pops, each O(log n): O(log n). find_median reads two roots: O(1). Space O(n) for all the numbers.

Why not choose the heap by size?

The tempting version is "push into whichever heap is shorter". That breaks order. Suppose lower = {1, 3} and upper = {5, 15}, and 20 arrives. Both heaps have size 2, so it goes to lower — and now lower holds 20, which is larger than everything in upper. Every median after that is wrong. The three-step version cannot make this mistake, because step 2 always passes the new number through the boundary before sizes are considered.

Edge cases

• One number: after step 3 it sits in lower, and the median is that number.
• Repeated values, such as five 7s: they split between the heaps; both roots are 7, and the median is 7.
• Negative numbers: lower stores -num, so a negative number becomes positive inside the heap. That is fine as long as you negate on the way out, which find_median does.
• Even count: the average must be a float. In Python 3 / always returns one; in Java, divide by 2.0.

Follow-ups

• "All numbers are between 0 and 100." Keep an array of 101 counts. Adding is O(1); finding the median walks at most 101 counts, which is O(1) too.
• "99% of numbers are between 0 and 100." Keep the 101 counts plus two counters for "below 0" and "above 100", with the rare outliers stored in small sorted lists. The median almost always falls in the counted range.
• "Median of a sliding window of size k" (Sliding Window Median). Two heaps again, but numbers must also leave. Use lazy deletion: record numbers to remove in a hash map and discard them when they reach a root.