AcWing 54. 数据流中的中位数

题目描述

如何得到一个数据流中的中位数？

如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。

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

样例

输入：1, 2, 3, 4

输出：1,1.5,2,2.5

解释：每当数据流读入一个数据，就进行一次判断并输出当前的中位数。

问题分析 

创建一个小顶堆和一个大顶堆，小顶堆保存比较大的那一半数，大顶堆保存比较小的那一半数，整个过程中都保证小顶堆的元素数量等于大顶堆的元素数量，或小顶堆比大顶堆多一个元素。当添加数时，先加到小顶堆里，然后将小顶堆的堆顶元素弹出加入到大顶堆里，这样就保证了添加数后，小顶堆中的数仍比大顶堆中的数大。然后判断如果小顶堆的元素数量小于大顶堆的元素数量，就把大顶堆的堆顶元素弹出加入到小顶堆中，这样就保证了小顶堆的元素数量大于等于大顶堆的元素数量。当要返回中位数时，判断如果大顶堆和小顶堆的元素数量相同，那么返回两个堆顶的平均值；如果小顶堆比大顶堆多一个元素，那么返回小顶堆的堆顶元素就是中位数。

代码实现

class Solution {
private:
    priority_queue<int> pq_max;
    priority_queue<int, vector<int>, greater<int>> pq_min;
    
public:
    void insert(int num){
        pq_min.push(num);
        pq_max.push(pq_min.top());
        pq_min.pop();
        if(pq_min.size() < pq_max.size()){
            pq_min.push(pq_max.top());
            pq_max.pop();
        }
    }

    double getMedian(){
        return pq_min.size() == pq_max.size()? 0.5 * (pq_min.top() + pq_max.top()) : pq_min.top();
    }
};