Leetcode 480. 滑动窗口中位数

Leetcode 480. 滑动窗口中位数

链接:原题传送门

--------------------

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

例如：

• [2,3,4]，中位数是 3
• [2,3]，中位数是 (2 + 3) / 2 = 2.5

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

--------------------

解法思路

k的奇偶性会影响中位数的计算方式，因此分开讨论。
先考虑k为奇数：
维护两个优先队列，一个size为（k/2+1）的大根堆q1，另一个size为(k/2)的小根堆q2。在滑动过程中，维持队列大小不变且保证q1的堆顶元素<=q2的堆顶元素。此时q1的堆顶元素即为中位数。
偶数：
滑动规则同奇数情况，但需要把q1的size固定为k/2，且中位数的计算方式为两个队列的堆顶元素的平均值。

需要注意的是STL的priority_queue不支持删除操作，所以引入了set来实时维护队列的大小。每次取堆顶元素时都判断一下该元素的下标是否在窗口内，如果不在窗口内，则直接跳过，直到取到一个下标在窗口内的堆顶元素。

class Solution {
public:


    struct CustomGreater {
        // pair, first 代表下标 ，second 代表元素值
        bool operator()(const std::pair<int, int>& a, const std::pair<int, int>& b) const {
                return a.second<b.second;
        }
    };


    struct CustomLess {
        bool operator()(const std::pair<int, int>& a, const std::pair<int, int>& b) const {
            return a.second > b.second;
        }
    };

    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        vector<int>kNums(k);
        vector<double>ans;
        priority_queue<pair<int,int>,vector<pair<int,int>>,CustomGreater>q1;
        priority_queue<pair<int,int>,vector<pair<int,int>>,CustomLess>q2;

        set<int>s1;
        int i=0;
        // 初始化第一个窗口
        for(;i<k;i++){
            q2.emplace(i,nums[i]);
        }

        for(int j=0;j<k/2+(k&1);j++){
            q1.push(q2.top());
            s1.insert(q2.top().first);
            q2.pop();
        }

        if(k&1)
            ans.push_back(q1.top().second);
        else
            ans.push_back((double(q1.top().second)+double(q2.top().second))/2);

        // 滑动
        for(;i<nums.size();i++){
            if(s1.contains(i-k)){
                // 左边队列的一个元素被移出窗口,要给左边补一个元素
                s1.erase(i-k);
                // 先把当前元素塞进右边队列
                q2.emplace(i,nums[i]);

                // 把右边队列的最小值拿出来 ，补进左边队列
                s1.insert(q2.top().first);
                q1.push(q2.top());
                q2.pop();
            }else{
                // 右边队列的某个元素被移除，需要补一个
                // 先把当前元素塞进左边队列
                q1.emplace(i,nums[i]);
                s1.insert(i);

                q2.push(q1.top());
                s1.erase(q1.top().first);
                q1.pop();

            }

            while(!q1.empty() && q1.top().first<=i-k)q1.pop();
            while(!q2.empty() && q2.top().first<=i-k)q2.pop();
            if(k&1)
                ans.push_back(q1.top().second);
            else
                ans.push_back((double(q1.top().second)+double(q2.top().second))/2);
        }

        return ans;
    }
};

pbds大法好

pbds库的平衡树提供find_by_order查询第k大，删除、插入、查询操作均为O(log k) 的时间复杂度。

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
class Solution {
public:
    tree<pair<int, int>, null_type, less<pair<int, int>>, rb_tree_tag, tree_order_statistics_node_update> t;
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        queue<pair<int, int>> q;
        vector<double> ans;
        for(int i = 0; i < nums.size(); ++i) {
            q.push({nums[i], i});
            t.insert({nums[i], i});
            if(q.size() >= k) {
                long x = (*t.find_by_order(k / 2)).first;
                long y = (*t.find_by_order((k - 1) / 2)).first;
                ans.push_back((x + y) / 2.0);
                t.erase(q.front());
                q.pop();
            }
        }
        return ans;
    }
};