C++实现一个数据流中,随时可以取得中位数

1. 题目描述

Leetcode295: 数据流的中位数
中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。

例如，

[2,3,4] 的中位数是 3

[2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构：

void addNum(int num) - 从数据流中添加一个整数到数据结构中。
double findMedian() - 返回目前所有元素的中位数。

示例：

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

2.思路分析

采用优先队列，也就是大根堆跟小根堆，我们将较小的n/2个数放到大根堆中，将较大的n/2个数放到小根堆中，显然，如果n是偶数，那么大根堆的堆顶跟小根堆的堆顶就是我们要找的两个中位数，将其相加除以2作为结果返回即可。如果n是奇数，那么就看大根堆跟小根堆谁的节点个数比另一个堆多一个，节点数量多的那个堆的堆顶就是我们要找的中位数，此时我们直接返回结果即可。

注意：

1. 对于取堆顶元素的操作的时间复杂度是常数级别的。

2. 插入新节点时我们需要判断节点的值是否小于大根堆堆顶的值或者大于小根堆堆顶的值，如果小于大根堆堆顶的值，那么节点应该插入大根堆，反过来应该插入小根堆。

3. 每次插入新节点我们还需要判断两个堆之间的元素个数是否平衡。插入新节点后，我们判断两个堆的元素个数，如果相差为2那么我们就要对堆进行调整。比如新插入一个节点到小根堆中，而此时大根堆的个数+1小于小根堆的节点个数，这个时候只需要将小根堆的堆顶元素弹出，然后将这个弹出的元素插入大根堆即可。反过来也是一样的操作。为什么可以这样做呢？这是因为我们说了小根堆保存的是较大的n/2个数，而小根堆的堆顶是小根堆中最小的元素，同时也是大根堆中最大的元素，因此我们将这个堆顶元素弹出并插入大根堆的操作并不会破坏“小根堆保存较大的n/2个数，大根堆保存较小的n/2”这样的前提。

3. 代码

这里还是借助了std::priority_queue来实现大跟堆和小根堆。

#include <iostream>
#include <vector>
#include <queue>

class MedianFinder {
 private:
     std::priority_queue<int, std::vector<int>, std::less<int> > maxHeap;
     std::priority_queue<int, std::vector<int>, std::greater<int> > minHeap;
 public:
    /** initialize your data structure here. */
    MedianFinder() {

    }

    void addNum(int num) {
        if (maxHeap.empty()) {  // if maxHeap is empty, push it directly
            maxHeap.push(num);
            return;
        }
        if (maxHeap.top()>= num) {  // small number push to maxHeap
            maxHeap.push(num);
        } else {    // max number push to minHeap
            if (minHeap.empty()) {  // if minHeap is empty, push it directly
                minHeap.push(num);
                return;
            }
            if (minHeap.top() >= num) { // push it to maxHeap
                maxHeap.push(num);
            } else {    // push it to minHeap
                minHeap.push(num);
            }
        }
        // adjust maxHeap and minHeap, ensure the different size <= 1
        if (maxHeap.size() == minHeap.size() + 2) { // maxHeap - minHeap = 2, adjust
            int maxTmp = maxHeap.top();
            minHeap.push(maxTmp);
            maxHeap.pop();
        }
        if (minHeap.size() == maxHeap.size() + 2) { // minHeap - maxHeap = 2, adjust
            int minTmp = minHeap.top();
            maxHeap.push(minTmp);
            minHeap.pop();
        }

    }

    double findMedian() {
        int maxHeapSize = maxHeap.size();
        int minHeapSize = minHeap.size();
        if (maxHeapSize + minHeapSize == 0) {   // empty return 0.0
            return 0.0;
        }
        int maxHeapTop = maxHeap.top();
        int minHeapTop = minHeap.top();
        if (((maxHeapSize + minHeapSize) & 1) == 0) {
            return (maxHeapTop + minHeapTop) / 2.0;
        }
        else {    // total sum is odd
            return maxHeapSize > minHeapSize ? (double)maxHeapTop : (double)minHeapTop;
        }
    }
};

 int main () {
    MedianFinder* obj = new MedianFinder();
    obj->addNum(1);
    obj->addNum(2);
    std::cout << obj->findMedian() << std::endl;
    obj->addNum(3);
    std::cout << obj->findMedian() << std::endl;

    return 0;
 }

4. 参考文献

1. LeetCode-295 数据流的中位数

贪心的题目暂时到这里，后面可能是动态规划的一些题目了。