[LeetCode 295] Find Median from Data Stream

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

For example,

[2,3,4], the median is 3

[2,3], the median is (2 + 3) / 2 = 2.5

Design a data structure that supports the following two operations:

• void addNum(int num) - Add a integer number from the data stream to the data structure.
• double findMedian() - Return the median of all elements so far.

Example:

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

Follow up:

1. If all integer numbers from the stream are between 0 and 100, how would you optimize it?
2. If 99% of all integer numbers from the stream are between 0 and 100, how would you optimize it?

分析

这道题的初一看想着使用vector来存放所有的num，然后add的时候使用二分法找到插入的点。findMedian时，直接找到中间的一个或者两个点即可获得结果。

但是这一题有一个更巧妙的解法，因为findMedian时只需要知道size和最中间的两个数即可，那么优先级队列能很好的解决这个问题。我们只需要使用一个最大优先级队列和一个最小优先级队列，在addNum的时候保持两个队列的size相等或者最大优先级队列size = 最小优先级队列size +1。那么getMedian的时候就可以通过队列的top()元素获取结果。

Code

class MedianFinder {
public:
    /** initialize your data structure here. */
    MedianFinder() {
        maxQ.push(INT_MIN);
        minQ.push(INT_MAX);
    }
    
    void addNum(int num) {
        int bigNum = minQ.top();
        int smallNum = maxQ.top();
        if (num >= smallNum)
        {
            minQ.push(num);
        }
        else
        {
            maxQ.push(num);
        }
        
        if (minQ.size() > maxQ.size() + 1)
        {
            int tmp = minQ.top();
            minQ.pop();
            maxQ.push(tmp);
        }
        else if (minQ.size() < maxQ.size())
        {
            int tmp = maxQ.top();
            maxQ.pop();
            minQ.push(tmp);
        }
        return;
        
    }
    
    double findMedian() {
        if ((minQ.size() + maxQ.size()) %2 == 1)
        {
            return minQ.top();
        }
        return (minQ.top() + maxQ.top()) / 2.0;
            
    }
    
    priority_queue<int> maxQ;
    priority_queue<int, vector<int>, greater<int>> minQ;
};

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder* obj = new MedianFinder();
 * obj->addNum(num);
 * double param_2 = obj->findMedian();
 */