Find Median from Data Stream (LeetCode 295)

The problem

Design a class MedianFinder with two methods:

• addNum(x) adds a number from the stream.
• findMedian() 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’s even.

Input:  addNum(1), addNum(2), findMedian(), addNum(3), findMedian()
Output: 1.5, 2

After 1 and 2 the middle two are 1 and 2, so the median is 1.5. After 3 the sorted stream is 1, 2, 3, and the middle value is 2.

How to think about it

Start with brute force. Keep every number in a list and sort it whenever the median is asked: O(n log n) per query. Keeping the list sorted with binary-search insertion makes queries O(1), but each insert still shifts up to n items, so O(n) per add. With many calls, that’s too slow.

What does the median actually need? Not the full order. It needs the largest number in the lower half and the smallest number in the upper half. Everything else can stay unordered.

The key insight: two heaps, one per half. Keep the lower half in a max-heap (low) so its largest value is at the root. Keep the upper half in a min-heap (high) so its smallest is at the root. Maintain two rules: every value in low ≤ every value in high, and low has the same size as high, or one more. Then the median is low’s root (odd count) or the average of both roots (even count), read in O(1).

Keeping the rules on each add. Push the new number into low, then move low’s largest over to high; that keeps the halves correctly ordered. If high is now bigger, move its smallest back. Each add is a few O(log n) heap operations. Which pattern is this? Two heaps. The tell is “running median”, or any question about the middle of a changing collection.

Approach

1. low is a max-heap (store negatives); high is a min-heap.
2. addNum(x): push x into low.
3. Pop low’s largest and push it into high. This keeps every low value ≤ every high value.
4. If high now has more items than low, move high’s smallest back into low.
5. findMedian(): if low is bigger, return its root; otherwise return the average of the two roots.

Solution

import heapq
class MedianFinder:
    def __init__(self):
        self.low = []    # max-heap (stored as negatives)
        self.high = []   # min-heap
    def addNum(self, x):
        heapq.heappush(self.low, -x)
        heapq.heappush(self.high, -heapq.heappop(self.low))   # balance value
        if len(self.high) > len(self.low):                    # balance size
            heapq.heappush(self.low, -heapq.heappop(self.high))
    def findMedian(self):
        if len(self.low) > len(self.high): return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2
# add 1,2 -> median 1.5 ; add 3 -> median 2

Dry run

Adding 5, 15, 1 and 3, with each heap shown sorted as positive values:

On the last add, 3 went into low, then low’s largest (5) moved up to high, so the halves stayed ordered and equal in size.

Complexity

• Time: O(log n) per addNum, O(1) per findMedian. An add does at most three heap pushes and two pops; reading the median only looks at the roots.
• Space: O(n). Every number is stored in one of the two heaps.

Edge cases and pitfalls

• Negative numbers work: after adding -1 and -2 the median is -1.5. Just be careful with the sign flips on low.
• Integer vs float. / 2 gives a float such as 10.0, which LeetCode accepts. In Java or C++, add as long or double to avoid overflow and integer division.
• Skipping the “balance value” step breaks the ordering. Pushing straight into whichever heap is smaller can leave a large value in low.
• Calling findMedian before any add would crash on low[0]. LeetCode never does this, but say so in an interview.
• Follow-up: if all numbers are in 0..100, a count array of size 101 answers the median in O(100) with no heaps.

Related problems

• Kth Largest Element in a Stream (LeetCode 703): a one-heap streaming statistic.
• Median of Two Sorted Arrays (LeetCode 4): the median of fixed sorted inputs, solved with binary search.
• Kth Largest Element in an Array (LeetCode 215): the min-heap-for-the-largest trick that high uses.