Find Median from Data Stream

The median is the middle value in a sorted list of integers. For lists of even length, there is no middle value, so the median is the mean of the two middle values.

For example:

• For arr = [1,2,3], the median is 2.
• For arr = [1,2], the median is (1 + 2) / 2 = 1.5

Implement the MedianFinder class:

• MedianFinder() initializes the MedianFinder object.
• void addNum(int num) adds the integer num from the data stream to the data structure.
• double findMedian() returns the median of all elements so far.

MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1); // arr = [1]
medianFinder.addNum(2); // arr = [1, 2]
medianFinder.findMedian(); // return 1.5
medianFinder.addNum(3); // arr = [1, 2, 3]
medianFinder.findMedian(); // return 2.0

Approach

Maintain two heaps: a max-heap for the lower half and a min-heap for the upper half. When adding a number, push it to the max-heap, move the max to the min-heap, then rebalance so the max-heap has either the same size or one extra element. The median is the top of the max-heap for odd counts or the average of both tops for even counts.

Why It Works

The max-heap stores the largest element of the lower half and the min-heap stores the smallest element of the upper half. Rebalancing keeps the halves sized correctly, so the median is always determined by the heap tops. This holds after every insertion.

Complexity

• Time: O(log n) per add, O(1) per median query
• Space: O(n)

Reference Implementation (Python)

class MedianFinder:

    def __init__(self):
        self.small = []  # max-heap (store negatives)
        self.large = []  # min-heap

    def addNum(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 findMedian(self) -> float:
        if len(self.small) > len(self.large):
            return float(-self.small[0])
        return (-self.small[0] + self.large[0]) / 2.0