中位数计算-大堆/小堆对顶

前言

        许多科学计算内容需要用到中位数的选取，例如KNN（K-近邻）算法，构造生成kd树的过程中需要逐个选择中位数以构造划分空间的超平面，或者在动态添加减少数据的过程中，依次获取中位数。获取中位数的一般做法是将数组排序，获取第N/2+1个元素（总数为奇数）或N/2与N/2+1个元素的平均值，对于一次普通查找中位数，这样操作的时间复杂度为排序时间复杂度O(logn)，但对于频繁插入、频繁减少数据、频繁查询中位数的操作而言，时间代价是很难承受的。这里介绍一个利用大顶堆和小顶堆进行平衡获取中位数的操作，可以方便数据插入和删除。

频繁添加数据与查询中位数

        以Leetcode中的题目295.数据流的中位数为例，简要解析数据流不断涌入时候的中位数获取方法。

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。
例如，
[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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/find-median-from-data-stream
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

        思考方法如下：将所有比中位数小的数值存在small堆中，比中位数大的数值存在big堆中，并且保证两堆容量之差小于等于1。这样，中位数就一定在两个堆的堆顶之中。我们将这种堆称为对顶堆。对顶堆长这样：

插入操作：

• 当插入数据大于small堆的堆顶之时，就将该数据插入到small堆中；否则插入到big堆中。
• 如果由于此次插入，使得两堆容量只差大于1，则调整两堆；将容量更多的那堆堆顶元素转到另一堆中。
  取中位数操作：
• 如果当前堆中元素总数为奇数，则一定出于两堆的堆顶元素之一，且为数字较多堆的堆顶元素。
• 如果当前总数为偶数，则中位数是两堆的堆顶元素平均值。

代码

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

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

滑动窗口获取中位数

        更进一步的，在数据中还需要做删除操作之后继续获取中位数，我们以Leetcode中的题目480. 滑动窗口中位数为例，简要解析数据流不断涌入时候且删除部分数据时候的中位数获取方法。

中位数是有序序列最中间的那个数。如果序列的长度是偶数，则没有最中间的数；此时中位数是最中间的两个数的平均数。

例如：

[2,3,4]，中位数是 3
[2,3]，中位数是 (2 + 3) / 2 = 2.5
给你一个数组 nums，有一个长度为 k 的窗口从最左端滑动到最右端。窗口中有 k 个数，每次窗口向右移动 1 位。你的任务是找出每次窗口移动后得到的新窗口中元素的中位数，并输出由它们组成的数组。

 

示例：

给出 nums = [1,3,-1,-3,5,3,6,7]，以及 k = 3。

窗口位置                      中位数
---------------               -----
[1  3  -1] -3  5  3  6  7       1
 1 [3  -1  -3] 5  3  6  7      -1
 1  3 [-1  -3  5] 3  6  7      -1
 1  3  -1 [-3  5  3] 6  7       3
 1  3  -1  -3 [5  3  6] 7       5
 1  3  -1  -3  5 [3  6  7]      6
 因此，返回该滑动窗口的中位数数组 [1,-1,-1,3,5,6]。

 

提示：

你可以假设 k 始终有效，即：k 始终小于等于输入的非空数组的元素个数。
与真实值误差在 10 ^ -5 以内的答案将被视作正确答案。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sliding-window-median
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

        本题考查动态维护数组的中位数。
        我们思考中位数的性质：如果一个数是中位数，那么在这个数组中，大于中位数的数目和小于中位数的数目，要么相等，要么就相差一。
        因此，我们采用对顶堆的做法，控制所有小于等于中位数的数字放到一个堆中，控制所有比中位数大的数字放到另一个堆中，并且保证两个堆的数目相差小于等于1。这样就可以保证每一次查询中位数的时候，答案一定出于两个堆的堆顶元素之一。
        因此选定数据结构：优先队列。因为优先队列采用的是堆结构，正好符合我们的需求。我们将所有小于等于中位数的元素放到small堆中（是一个大顶堆），将所有大于中位数的元素放到big堆中（是一个小顶堆）。

        初始化方法如下：

• 将前K个元素全部插入到small堆中。从small堆中弹出K/2个元素到big堆中。
• 这样，当K为奇数，则small堆元素比big堆元素多1；当K为偶数，两个堆元素相等。

        取中位数的操作：

• 我们的插入操作可以保证每次优先插入到small堆中，因此small堆中的元素个数大于等于big堆的元素个数。
• 当K为奇数时候，中位数是元素数量较多的small堆 堆顶元素。
• 当K为偶数时候，中位数是small堆和big堆的堆顶元素平均值。

        窗口滑动过程中的操作：

• 假定在上一次滑动之后，已经有small堆和big堆元素数目相差小于等于1.
• 设置当前的滑动时，balance = 0。balance表示small堆元素数目减去big堆元素个数。
• 删除窗口左侧的元素。
• • 由于堆无法直接删除掉某个指定元素，先欠下这个账，等某次元素出现在堆顶的时候，再删除他。mp记录这个元素欠账的个数。mp[left]++;
• • 虽然没有真的在堆数据结构中删除窗口最左侧的元素，但是在我们的心中已经删掉他了。堆两侧的平衡性发生了变化。如果left<=small.top()，就说明删掉的元素在small堆中，我们让balance–;否则，就说明删掉的元素在big堆中，让balance++;
• 添加进来窗口右侧的元素。如果right<=small.top()，就应该让这个元素放到samll堆里面，balance++；否则放到big堆里，balance–。
• 经过上面的操作，balance要么为0，要么为2，要么为-2。我们需要经过调整使得balance为0。
• • 如果balance为0，在这次窗口滑动之前已经是平衡的，这次调整也没有让两堆的数目变化，所以不用调整两边的堆。
• • 如果balance为2，就说明small堆的元素比big堆的元素多了两个。从small堆减少一个，big堆里增加一个，就可以让两边相等。big.push(small.top());small.pop();
• • 如果balance为-2，就说明big堆的元素比small堆的元素多了两个。从big堆减少一个，small堆里增加一个，就可以让两边相等。small.push(big.top());big.pop();
• 调整完了，现在该欠债还钱了。不能让那些早该删除的元素涉及到中位数的运算。
• • 分别检查两边的堆顶元素，如果堆顶元素欠着债，则弹出堆顶元素，直到堆顶元素没有欠债为止。有朋友问了：堆顶下面也有欠债的怎么办呢？我们之前说过，取中位数的时候只与堆顶元素有关，至于那些堆顶下面欠着债的，欠着就欠着吧，等他们到堆顶的时候再弹出去就好了。
• 最后，添加中位数即可。

代码

class Solution {
public:
    priority_queue<int> small;
    priority_queue<int, vector<int>, greater<int> > big;
    unordered_map<int, int> mp;
    double get(int& k){
        if(k%2) return small.top();
        else return ((long long)small.top()+big.top())*0.5;
    }
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        for(int i = 0; i < k; i++){small.push(nums[i]);};
        for(int i = 0; i < k / 2; i++){big.push(small.top()); small.pop();}
        vector<double> ans{get(k)};
        for(int i = k; i < nums.size(); i++){
            int balance = 0;
            int l = nums[i-k];
            mp[l]++;
            if(!small.empty() && l<=small.top()){balance--;}
            else {balance++;}
            if(!small.empty() && nums[i] <= small.top()){
                small.push(nums[i]);
                balance++;
            }
            else{
                big.push(nums[i]);
                balance--;
            }
            if(balance>0){
                big.push(small.top());
                small.pop();
            }
            if(balance<0){
                small.push(big.top());
                big.pop();
            }
            while(!small.empty() && mp[small.top()]>0){
                mp[small.top()]--;
                small.pop();
            }
            while(!big.empty() && mp[big.top()]>0){
                mp[big.top()]--;
                big.pop();
            }
            ans.push_back(get(k));
            
        }
        return ans;
    }
};

最后附上我的leetcode主页ProfSnail，欢迎大家共同讨论、共同学习。