roadmap/heap/Find Median from Data Stream

Find Median from Data Stream

HardAmazonGoogleMetaMicrosoft

Problem

Design a data structure that supports addNum(int num) and findMedian() → return the median of all added numbers.

Intuition — the key insight

Split numbers into two halves: a max-heap for the lower half and a min-heap for the upper half. Keep them balanced (sizes differ by at most 1). Median is max of lower heap (odd total) or average of both tops (even total).

Approach

Two heaps

smallHeap (max-heap, lower half), largeHeap (min-heap, upper half). On add: push to smallHeap, then balance by moving small's max to largeHeap if needed, then re-balance sizes.

Complexity

Time

O(log n) add, O(1) findMedian

Space

O(n)

Heap operations are O(log n). Peek is O(1).

Solution code

Python

class MedianFinder:
    def __init__(self):
        self.small = []  # max-heap (negate values)
        self.large = []  # min-heap
    
    def addNum(self, num: int) -> None:
        heappush(self.small, -num)
        # Ensure all small <= all large
        if self.small and self.large and (-self.small[0] > self.large[0]):
            heappush(self.large, -heappop(self.small))
        # Balance sizes
        if len(self.small) > len(self.large) + 1:
            heappush(self.large, -heappop(self.small))
        if len(self.large) > len(self.small):
            heappush(self.small, -heappop(self.large))
    
    def findMedian(self) -> float:
        if len(self.small) > len(self.large):
            return -self.small[0]
        return (-self.small[0] + self.large[0]) / 2

Edge cases to remember

⚠Single element: return it

⚠Even count: average of two middle elements

⚠All same values: median is that value

Hints

What pattern do you think this problem uses?