Categories
不学无术

LeetCode 295. Find Median from Data Stream

题目大意是给定一个int数据流,判断当前整个数组的中位数值。由于数据流是可以不断增加的,因此采用两个堆(最大堆*1+最小堆*1)来实现。
这东西思路网上都有,每次新进一个数,维护堆就行了。维护一次的复杂度是O(logN),对应整个数组的复杂度是O(N logN)

STL 堆相关的三个函数 make_heap(), push_heap(), pop_heap()

STL是个好东西,在<algorithm>库里面就有三个与堆有关的函数,利用vector存储堆元素就可以调用这些函数来建立、更新堆。
make_heap(first, last, comp)是建立堆的函数,输入的first, last对应begin()和end()的iterator,comp函数可选,与库中sort()用到的方法相同,自定义的时候比较两个元素x, y的偏序关系,返回true说明x应该排在前面。相关的有两个仿函数叫less<T>() 和greater<T>(),直接用来比较<和>关系的。
push_heap(first, last, comp)使用时,将新元素放在数组末尾(代表堆的末叶节点),然后调用这个函数更新堆,是一个sift-up的过程。
pop_heap(first, last, comp)会将堆顶元素调整到数组末尾用于pop,同时调整更新堆。

代码

class MedianFinder {
private:
    vector<int> maxHeap; // store values less than median
    vector<int> minHeap; // store values greater than median
public:
    // Adds a number into the data structure.
    void addNum(int num) {
        int minSize = minHeap.size();
        int maxSize = maxHeap.size();
        if(maxSize == 0){
            maxHeap.push_back(num);
            return;
        }
        if(num < maxHeap[0]){
            maxHeap.push_back(num);
            push_heap(maxHeap.begin(), maxHeap.end());
            if(maxSize > minSize) {
                pop_heap(maxHeap.begin(), maxHeap.end());
                int popVal = maxHeap[maxHeap.size()-1];
                maxHeap.pop_back();
                minHeap.push_back(popVal);
                push_heap(minHeap.begin(), minHeap.end(), greater<int>());
            }
        } else {
            minHeap.push_back(num);
            push_heap(minHeap.begin(), minHeap.end(), greater<int>());
            if(minSize >= maxSize) {
                pop_heap(minHeap.begin(), minHeap.end(), greater<int>());
                int popVal = minHeap[minHeap.size()-1];
                minHeap.pop_back();
                maxHeap.push_back(popVal);
                push_heap(maxHeap.begin(), maxHeap.end());
            }
        }
    }
    // Returns the median of current data stream
    double findMedian() {
        if((maxHeap.size() + minHeap.size())%2 == 0){
            //even num.
            if(maxHeap.empty() && minHeap.empty())
                return 0.0;
            double sum = maxHeap[0] + minHeap[0];
            return sum/2.0;
        } else {
            if(maxHeap.empty())
                return 0.0;
            return (double)maxHeap[0];
        }
    }
};
// Your MedianFinder object will be instantiated and called as such:
// MedianFinder mf;
// mf.addNum(1);
// mf.findMedian();