Find Median From Data Stream

Problem

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.

Examples

Example 1:

Input:
["MedianFinder", "addNum", "1", "findMedian", "addNum", "3" "findMedian", "addNum", "2", "findMedian"]

Output:
[null, null, 1.0, null, 2.0, null, 2.0]

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

Constraints

• -100,000 <= num <= 100,000
• findMedian will only be called after adding at least one integer to the data structure.

You should aim for a solution with O(log n) time for addNum(), O(1) time for findMedian(), and O(n) space, where n is the current number of elements.

Solution

class MedianFinder {
    std::priority_queue<int, std::vector<int>, std::less<int>> max_heap;
    std::priority_queue<int, std::vector<int>, std::greater<int>> min_heap;
public:
    MedianFinder() = default;
    
    void addNum(int num) {
        // If num is greater than min_heap top, add to min heap
        if (!min_heap.empty() && num > min_heap.top()) {
            min_heap.push(num);
        } else {
            max_heap.push(num);
        }
        // Rebalance if size difference is larger than 1
        auto max_size = std::max(min_heap.size(), max_heap.size());
        auto min_size = std::min(min_heap.size(), max_heap.size()
        if (max_size - min_size) > 1) {
            if (min_heap.size() > max_heap.size()) {
                max_heap.push(min_heap.top());
                min_heap.pop();
            } else {
                min_heap.push(max_heap.top());
                max_heap.pop();
            }
        }
    }
    
    double findMedian() {
        if (min_heap.size() == max_heap.size()) {
            return static_cast<double>(min_heap.top() + max_heap.top()) / 2.0;
        }
        return static_cast<double>(
            min_heap.size() > max_heap.size() ? min_heap.top() : max_heap.top()
        );
    }
};