刷题： 数据流中的中位数

持续创作，加速成长！这是我参与「掘金日新计划 · 6 月更文挑战」的第24天，点击查看活动详情

题目来源：剑指 Offer 41. 数据流中的中位数

题目：

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

例如，

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

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

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

void addNum(int num) - 从数据流中添加一个整数到数据结构中。

double findMedian() - 返回目前所有元素的中位数。

示例 1：

输入：["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]

[[],[1],[2],[],[3],[]]

输出：[null,null,null,1.50000,null,2.00000]

示例 2：

输入：["MedianFinder","addNum","findMedian","addNum","findMedian"]

[[],[2],[],[3],[]]

输出：[null,null,2.00000,null,2.50000]  

限制：

最多会对 addNum、findMedian 进行 50000 次调用。

思路分析：

题目要求得到一个数据流排序后的中位数，数据流是动态的，且数据量可能会比较大，因此对数据流排序后求中位数是不可行的。

要求中位数的话，那么分别用两个堆：大根堆用于记录小于等于中位数的数字，小根堆用于记录大于中位数的数字。且规定：当元素总数为偶数时，大根堆和小跟堆中的元素个数相等，当元素总数为奇数时，大根堆的元素个数比小跟堆多一个。

这样一来，如果元素总数为奇数，那么答案就是 大根堆 的堆顶元素；如果元素总数为偶数，那么答案就是大根堆和小跟堆的堆顶元素的平均值。

AC代码：

class MedianFinder {
public:
    priority_queue<int> small;
    priority_queue<int, vector<int>, greater<int> > big;
    int n;
    MedianFinder() {
        n = 0;
    }
    void addNum(int num) {
        if(small.empty()){
            small.push(num);
            n++;
            return;
        }
        if(num<=small.top()){
            small.push(num);
            n++;
        }
        else{
            big.push(num);
            n++;
        }
        if(small.size()-big.size()==2){
            big.push(small.top());
            small.pop();
        }
        if(big.size()-small.size()==2){
            small.push(big.top());
            big.pop();
        }
    }
    
    double findMedian() {
        if(n%2){
            if(small.size()>big.size())
                return small.top();
            return big.top();
        }
        else{
            return ((long long)small.top() + big.top()) * 0.5;
        }
    }
};