【LeetCode-480】滑动窗口中位数

问题

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

示例

输入： nums = [1,3,-1,-3,5,3,6,7], k = 3
输出： [1,-1,-1,3,5,6]

解答

class Solution {
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        vector<double> res;
        multiset<int> small, big; // 分别设为大根堆和小根堆
        for (int i = 0; i < nums.size(); i++) {
            if (small.size() > big.size()) {
                small.insert(nums[i]);
                big.insert(*small.rbegin());
                small.erase(prev(small.end()));
            } else {
                big.insert(nums[i]);
                small.insert(*big.begin());
                big.erase(big.begin());
            }
            if (i >= k) { // 执行删除操作
                auto spos = small.find(nums[i - k]);
                if (spos != small.end()) small.erase(spos);// 要删除的在大根堆中
                else big.erase(big.find(nums[i - k])); // 否则肯定在小根堆中
                if (small.size() + 1 == big.size()) { // 恢复堆的平衡关系
                    small.insert(*big.begin());
                    big.erase(big.begin());
                }
                if (small.size() - 2 == big.size()) { // 恢复堆的平衡关系
                    big.insert(*small.rbegin());
                    small.erase(prev(small.end()));
                }
            }
            if (i >= k - 1) { // 执行加入res操作
                if (k & 1) res.push_back(*small.rbegin());
                else res.push_back(((double)*small.rbegin() + *big.begin()) / 2);
            }
        }
        return res;
    }
};

重点思路

本题可以看作【剑指Offer-59-I】滑动窗口的最大值和【剑指Offer-41】数据流中的中位数的结合与拓展。

本题可以看作求可以删除元素的数据流的中位数。但是c++中的优先队列不支持删除任意元素，所以本算法使用multiset替代priority_queue实现大小顶堆的功能。multiset是有序的，begin()代表最小值，rbegin()代表最大值，prev(end())代表最大值迭代器所在的位置（这个需要特别注意，删除元素时会用到）。

本题的重难点在于删除操作中，如何在两个堆中找出我们要删除的值，以及删除后如何保持两个堆的平衡。由于数字是有重复的，我们要删除的数可能会存在于两个堆中。我们不需要特别去判断该数是两个堆中的哪一个，直接随便删除一个，然后再调整两个堆的平衡即可。

设小根堆为big，大根堆为small，“两个堆平衡”的定义为：big.size() == small.size()或者big.size() + 1 == small.size()。我们只需要讨论所有情况，设定对应的策略即可。