【动态维护前 x 大元素】LeetCode 3321. 计算子数组的 x-sum II

题目

https://leetcode.cn/problems/find-x-sum-of-all-k-long-subarrays-ii/description/

题解

定义两个有序集合 \(L, R\) 动态维护数组前 \(x\) 大元素。以滑动窗口的方式遍历数组，窗口大小为 \(n - k + 1\)，其中有序集合 \(R\) 用于维护当前窗口前 \(x\) 大的元素（可能不足 \(x\) 个），有序集合 \(L\) 维护当前窗口剩余的其它元素。

参考代码

using PLI = pair<long long, int>;

class Solution {
private:
    set<PLI> L, R;
    long long res = 0LL;

    void add(PLI p, int x) {
        if (R.size() < x) {
            R.insert(p);
            res += p.first * p.second;
        } else if (*R.begin() > p) {
            L.insert(p);
        } else {
            PLI d = *R.begin();
            L.insert(d);
            res -= d.first * d.second;
            R.erase(R.begin());
            R.insert(p);
            res += p.first * p.second;
        }
    }

    void del(PLI p) {
        if (R.find(p) != R.end()) {
            R.erase(p);
            res -= p.first * p.second;
            if (!L.empty()) {
                PLI a = *L.rbegin();
                res += a.first * a.second;
                R.insert(a);
                L.erase(a);
            }
        } else {
            L.erase(p);
        }
    }

public:
    vector<long long> findXSum(vector<int>& nums, int k, int x) {
        int n = nums.size();
        vector<long long> ans(n - k + 1);
        unordered_map<int, int> mp;
        for (int i = 0; i < n; ++ i) {
            if (i >= k) {
                del({mp[nums[i - k]], nums[i - k]});
                if (-- mp[nums[i - k]] > 0) {
                    add({mp[nums[i - k]], nums[i - k]}, x);
                } 
            }
            if (mp[nums[i]] > 0) {
                del({mp[nums[i]], nums[i]});
            }
            add({++ mp[nums[i]], nums[i]}, x);
            if (i >= k - 1) {
                ans[i - k + 1] = res;
            }
        }
        return ans;
    }
};