Leetcode 295. Find Median from Data Stream

Problem:

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?

 

Solution:

　　这道题要我们在数据流中找出中值，用数组的话插入的时间复杂度就是线性的，用set或二叉搜索树又不能做到随机查询，看了这道题的标签有个Heap我才知道用优先级队列。这里用一个大顶堆small和小顶堆large，大顶堆里是小于中位数的所有数字，小顶堆里是大于中位数的所有数字，并且我们确保大顶堆的元素个数只比小顶堆多1或相等。这样如果两个堆元素个数相等，就取堆顶取平均，否则取small的堆顶。对于follow up1，我们可以使用counting sort，用一个数组记录0到100元素出现的个数，每次插入把计数器加1，这样是O(1)的插入效率，查询时只需找到中间值在哪个元素上即可，查询效率也是O(1)。对于follow up2，我们则使用counting sort加优先级队列即可，是原问题和follow up1的结合。

Code:

 

class MedianFinder {
public:
    struct mycompare{
        bool operator()(int &a,int &b){
            return a>b;
        }
    };
    /** initialize your data structure here. */
    MedianFinder() {
        
    }
    
    void addNum(int num) {
        if(small.empty()){
            small.push(num);
            return;
        }
        if(small.size() == large.size()){
            if(num >= large.top()){
                small.push(large.top());
                large.pop();
                large.push(num);
            }
            else
                small.push(num);
        }
        else{
            if(num <= small.top()){
                large.push(small.top());
                small.pop();
                small.push(num);
            }
            else
                large.push(num);
        }
    }
    
    double findMedian() {
        if(small.size() == large.size())
            return (small.top()+large.top())/2.0;
        else
            return small.top();
    }
private:
    priority_queue<int> small;
    priority_queue<int,vector<int>,mycompare> large;
};

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