Find Median from Data Stream

Problem Description

The median is the middle value in an ordered integer list. If the size of the list is even, there is no single middle value, and the median is the average of the two middle values.

Implement the MedianFinder class:

• MedianFinder() Initializes the object.
• void addNum(int num) Adds the integer num to the data structure.
• double findMedian() Returns the median of all elements so far.

Examples

Constraints

• -10⁵ ≤ num ≤ 10⁵
• There will be at least one element in the data structure before calling findMedian.
• At most 5 * 10⁴ calls will be made to addNum and findMedian.

Dividing the Stream into Two Balanced Halves

The simple approach is maintaining a sorted array. To insert a number, we find its position using binary search and shift elements. This makes addNum O(n) due to element shifting, which is too slow for large streams.

To optimize, notice that we only need to access the middle elements. We do not need the entire list sorted. We can split the numbers into two halves: the smaller half and the larger half.

• Store the smaller half in a max-heap (so we can quickly access the largest of the small numbers).
• Store the larger half in a min-heap (so we can quickly access the smallest of the large numbers).

By keeping the heaps balanced (their sizes differ by at most one), the median is always either the top of the larger heap, or the average of the tops of both heaps.

Solution 1: Insertion Sort (Baseline)

Maintain a sorted array by inserting each incoming number at its correct position using binary search.

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

class MedianFinder {
    private List<Integer> list = new ArrayList<>();

    public void addNum(int num) {
        int idx = Collections.binarySearch(list, num);
        if (idx < 0) idx = -(idx + 1);
        list.add(idx, num); // shifts elements, costing O(n)
    }

    public double findMedian() {
        int n = list.size();
        if (n % 2 == 1) return list.get(n / 2);
        return (list.get(n / 2 - 1) + list.get(n / 2)) / 2.0;
    }
}

import bisect


class MedianFinder:
    def __init__(self):
        self.arr = []

    def addNum(self, num: int) -> None:
        # binary search + insert
        bisect.insort(self.arr, num)

    def findMedian(self) -> float:
        n = len(self.arr)
        if n % 2 == 1:
            return float(self.arr[n // 2])
        return (self.arr[n // 2 - 1] + self.arr[n // 2]) / 2.0

#include <vector>
#include <algorithm>

class MedianFinder {
    std::vector<int> arr;
public:
    void addNum(int num) {
        auto it = std::lower_bound(arr.begin(), arr.end(), num);
        arr.insert(it, num); // shifts elements, costing O(n)
    }

    double findMedian() {
        int n = arr.size();
        if (n % 2 == 1) return arr[n / 2];
        return (arr[n / 2 - 1] + arr[n / 2]) / 2.0;
    }
};

Complexity Analysis:

• Time Complexity: addNum is O(n) due to elements shift, findMedian is O(1).
• Space Complexity: O(n) to store the stream elements.

Where it breaks: If the stream has millions of elements, shifting arrays on every insertion causes unacceptable latency.

Solution 2: Two Heaps (Optimal)

Maintain a max-heap for the smaller half and a min-heap for the larger half.

import java.util.Collections;
import java.util.PriorityQueue;

class MedianFinder {
    private PriorityQueue<Integer> small = new PriorityQueue<>(Collections.reverseOrder()); // max-heap
    private PriorityQueue<Integer> large = new PriorityQueue<>(); // min-heap

    public void addNum(int num) {
        small.offer(num);
        // ensure all elements in small are <= elements in large
        large.offer(small.poll());

        // balance sizes (small can have at most 1 more element than large)
        if (small.size() < large.size()) {
            small.offer(large.poll());
        }
    }

    public double findMedian() {
        if (small.size() > large.size()) {
            return small.peek();
        }
        return (small.peek() + large.peek()) / 2.0;
    }
}

import heapq


class MedianFinder:
    def __init__(self):
        self.small = []  # max-heap (we store negative numbers to simulate max-heap)
        self.large = []  # min-heap

    def addNum(self, num: int) -> None:
        # push to small, then move the largest to large
        heapq.heappush(self.small, -num)
        val = -heapq.heappop(self.small)
        heapq.heappush(self.large, val)

        # balance sizes
        if len(self.small) < len(self.large):
            val = heapq.heappop(self.large)
            heapq.heappush(self.small, -val)

    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

#include <queue>
#include <vector>

class MedianFinder {
    std::priority_queue<int> small; // max-heap
    std::priority_queue<int, std::vector<int>, std::greater<int>> large; // min-heap
public:
    void addNum(int num) {
        small.push(num);
        large.push(small.top());
        small.pop();

        if (small.size() < large.size()) {
            small.push(large.top());
            large.pop();
        }
    }

    double findMedian() {
        if (small.size() > large.size()) {
            return small.top();
        }
        return (small.top() + large.top()) / 2.0;
    }
};

Complexity Analysis:

• Time Complexity: addNum is O(log n) because heap push/pop operations take logarithmic time. findMedian is O(1) as we only peek the heap roots.
• Space Complexity: O(n) to store the stream elements in both heaps.

Where it breaks: If there are float values in the stream, this specific integer implementation must be updated to float or double types. The logic remains identical.

Common Mistakes

• Forgetting that Python’s heapq is a min-heap. You must negate values when pushing to the max-heap, and negate them again when popping to restore the original value.
• Not maintaining heap balance. If you fail to rebalance the heap sizes, the peaks will not correspond to the true middle values of the stream.
• Integer division instead of floating point division. When dividing by 2 to compute the median of an even number of elements, ensure you use 2.0 (or cast to double) to avoid losing the decimal component.

Frequently Asked Questions

Can we solve this if the numbers are restricted to a specific range, say 0 to 100? Yes. You can use bucket sort (an array of counters for each number). addNum increments the counter in O(1) time. findMedian scans the counters to find the middle indices in O(100) = O(1) time. This is more optimal than the heap solution.

What is the benefit of the two-heap approach over binary search trees? Self-balancing binary search trees (like AVL trees) also support O(log n) insertion and O(1) median queries, but they have significant constant-factor overhead due to node pointers and rotation checks. Heaps are array-backed and perform much faster in practice.

What does this problem test in interviews? It tests your capability to design data structures, understand heap mechanics, and apply balancing algorithms to dynamic sorting tasks.

← All Problems