Find Median from Data Stream

How it works

To efficiently find the median while numbers keep coming, we split the stream into two halves:

• A max-heap (small) that stores the smaller half of the numbers.
  • The largest number of this half is on top.
• A min-heap (large) that stores the larger half of the numbers.
  • The smallest number of this half is on top.

The goal:

• Ensure both heaps are balanced in size (difference at most 1).
• Ensure all numbers in small are ≤ all numbers in large.

This setup allows:

• Median = top of the bigger heap (if odd count)
• Median = average of both tops (if even count)

This gives O(log n) insert and O(1) median lookup.

Steps

1. Initialize

   • Create two heaps:
     • small -> max-heap for lower half
     • large -> min-heap for upper half

2. addNum(x)

   • If large is not empty and x is greater than the smallest element in large, insert into large.
   • Otherwise insert into small.
   • Rebalance:
     • If one heap becomes larger than the other by more than 1, move the top element to the other heap.

3. findMedian()

   • If one heap has more elements:
     • Median = top of that heap.
   • If both heaps have equal elements:
     • Median = average of both heap tops.

Complexity

• Time complexity: O(m * log n) for addNum(), O(m) for findMedian().
• Space complexity: O(n)

Where m is the number of function calls and n is the length of the array.

Things to watch for

Using Wrong Heap Types

The two-heap solution requires a max-heap for the smaller half and a min-heap for the larger half. Swapping these or using two min-heaps produces incorrect medians. In languages like Python where heapq only provides min-heaps, you must negate values to simulate a max-heap for the smaller half.

Failing to Maintain Heap Balance

The heaps must differ in size by at most one element. Forgetting to rebalance after insertions leads to incorrect median calculations. After every addNum call, check if one heap has more than one extra element and transfer the top element to the other heap.

Incorrect Median Calculation for Even Count

When both heaps have equal size, the median is the average of both tops, not just one of them. Returning only the top of one heap or using integer division instead of floating-point division produces wrong results. Always check heap sizes and compute (smallTop + largeTop) / 2.0 for the even case.

Integer Overflow in Sum Calculation

When computing the average of two heap tops, adding two large integers can cause overflow in languages like Java or C++. Cast to a larger type before adding, or compute as a + (b - a) / 2.0 to avoid overflow while still getting the correct floating-point result.

Inserting Into Wrong Heap Initially

The first element must go into one of the heaps, but subsequent elements must be compared against the appropriate heap top to determine placement. Inserting all elements into one heap first and then rebalancing works, but directly inserting into the wrong heap without proper comparison breaks the invariant that all elements in the small heap are less than or equal to all elements in the large heap.

import heapq

class MedianFinder:
    def __init__(self):
        # two heaps, large, small, minheap, maxheap
        # heaps should be equal size
        self.small, self.large = [], []

    def addNum(self, num: int) -> None:
        if self.large and num > self.large[0]:
            heapq.heappush(self.large, num)
        else:
            heapq.heappush(self.small, -1 * num)

        if len(self.small) > len(self.large) + 1:
            val = -1 * heapq.heappop(self.small)
            heapq.heappush(self.large, val)
        if len(self.large) > len(self.small) + 1:
            val = heapq.heappop(self.large)
            heapq.heappush(self.small, -1 * val)

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

class MedianFinder {
    priority_queue<int, vector<int>, less<int>> smallHeap;
    priority_queue<int, vector<int>, greater<int>> largeHeap;

public:
    MedianFinder() {}

    void addNum(int num) {
        smallHeap.push(num);
        if (!largeHeap.empty() && smallHeap.top() > largeHeap.top()) {
            largeHeap.push(smallHeap.top());
            smallHeap.pop();
        }
        if (smallHeap.size() > largeHeap.size() + 1) {
            largeHeap.push(smallHeap.top());
            smallHeap.pop();
        }
        if (largeHeap.size() > smallHeap.size() + 1) {
            smallHeap.push(largeHeap.top());
            largeHeap.pop();
        }
    }

    double findMedian() {
        if (smallHeap.size() == largeHeap.size()) {
            return (largeHeap.top() + smallHeap.top()) / 2.0;
        } else if (smallHeap.size() > largeHeap.size()) {
            return smallHeap.top();
        } else {
            return largeHeap.top();
        }
    }
};