【动态维护前 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;}
};