C++ 随笔：用两个有序集合维护滑动窗口
在需要动态维护滑动窗口内的统计量（如前 K 小值之和或中位数）时，可以使用两个有序集合来实现高效更新。
基本思路：
先确定窗口长度以及在窗口中要维护的元素数量，然后用两个集合 big/small 来区分管理：big 存储当前窗口中符合要求的 K 个元素；small 存储窗口中的其他元素；根据题意决定使用 set 或 multiset，前者不允许重复，后者可以处理重复值。
rebalance
在维护过程中，关键是 rebalance 操作。它负责保持两个集合之间的数量与顺序约束，当 big 的元素数少于目标数量时，就从 small 中取出最高优先级元素转移到 big；若 small 中存在比 big 中优先级更高的元素，则交换两者，使得 big 始终保存窗口中优先级最高的那部分。
auto rebalance=[&](){
// 数量少则补
while(big.size()<k-1 && !small.empty()){
auto it=small.begin();
sum+=*it;
big.insert(*it);
small.erase(it);
}
// 优先级低则换
while(!big.empty() && !small.empty()){
auto r = prev(big.end()),l=small.begin();
int R=*r, L=*l;
if(L<R){
sum+=L-R;
big.erase(r);
big.insert(L);
small.erase(l);
small.insert(R);
}else break;
}
};
add/remove
更新时，add 和 remove 分别对应窗口的右移操作。add 用于加入新进入窗口的元素，先插入到 small，再执行一次 rebalance；remove 用于移除滑出窗口的元素，优先从 big 中删除，若不存在再在 small 中删除，之后同样调用 rebalance；根据题意/元素类型看是否添加其他逻辑。
auto add=[&](int k){
// 最基础的add，根据题意是否添加其他逻辑
small.insert(k);
rebalance();
};
auto erase_one=[&](int item){
// 删除元素item，优先从 big 中删除
if(auto it=big.find(item); it!=big.end()){
sum-=item;
big.erase(it);
}else if(auto it=small.find(item); it!=small.end()){
small.erase(it);
}
};
auto remove=[&](int k){
// 最基础的remove，根据题意是否添加其他逻辑
erase_one(k);
rebalance();
};
这种两个有序集合维护滑动窗口的技巧，本质上是通过平衡集合间的划分来动态维护一个有序的前缀
问题举例
总体上这类问题的思路就是这样，下面取 LeetCode 中的几个例子：
LeetCode 3321
首先是将原集合中对应的元素删除，之后再将key对应的value操作：加一 or 减一，加一蛮简单的，减一要注意判断边界；如果value归零了，那么就后续不需要在insert进small_set中，否则都需要insert进small_set；最后再通过一个rebalance操作，主要目的将 big_set 补全，数量不足从small中补，比较big当前最小值与small最大值，符合则替换；
struct cmp{
bool operator()(const pair<int,int>& a,const pair<int,int>& b)const{
if(a.first==b.first)return a.second>b.second;
return a.first>b.first;
}
};
vector<long long> findXSum(vector<int>& nums, int k, int x){
unordered_map<int,int>q;
set<pair<int,int>,cmp>big,small;
long long sum=0;
auto rebalance=[&](){
while(big.size()<x && !small.empty()){
auto it=small.begin();
sum+=1ll*it->first*it->second;
big.insert(*it);
small.erase(it);
}
while(!big.empty() && !small.empty()){
auto r = prev(big.end()),l=small.begin();
if(cmp()(*l,*r)){
sum+=1ll*l->first*l->second-1ll*r->first*r->second;
auto R=*r,L=*l;
big.erase(r);
big.insert(L);
small.erase(l);
small.insert(R);
}else break;
}
};
auto erase_one=[&](pair<int,int>item){
if(auto it=big.find(item); it!=big.end()){
sum-=1ll*it->first*it->second;
big.erase(it);
}else if(auto it=small.find(item); it!=small.end()){
small.erase(it);
}
};
auto add=[&](int k){
int v=q[k];
erase_one({v,k});
q[k]=v+1;
small.insert({v+1,k});
rebalance();
};
auto remove=[&](int k){
int v=q[k];
erase_one({v,k});
q[k]=v-1;
if(q[k])small.insert({v-1,k});
rebalance();
};
for(int i=0;i<k;i++)add(nums[i]);
vector<long long>res{sum};
for(int i=k;i<nums.size();i++){
remove(nums[i-k]);
add(nums[i]);
res.push_back(sum);
}
return res;
}
LeetCode 3013
这样例子要简单一些，就是维护滑动窗口中 k-1 个最小值，可能存在的问题是 滑动窗口的长度和其中要维护的元素个数确定，这点要留心题意，之后的就是固定模板了
long long minimumCost(vector<int>& nums, int k, int x){
multiset<int>big,small;
long long sum=nums[0];
auto rebalance=[&](){
while(big.size()<k-1 && !small.empty()){
auto it=small.begin();
sum+=*it;
big.insert(*it);
small.erase(it);
}
while(!big.empty() && !small.empty()){
auto r = prev(big.end()),l=small.begin();
int R=*r, L=*l;
if(L<R){
sum+=L-R;
big.erase(r);
big.insert(L);
small.erase(l);
small.insert(R);
}else break;
}
};
auto erase_one=[&](int item){
if(auto it=big.find(item); it!=big.end()){
sum-=item;
big.erase(it);
}else if(auto it=small.find(item); it!=small.end()){
small.erase(it);
}
};
auto add=[&](int k){
small.insert(k);
rebalance();
};
auto remove=[&](int k){
erase_one(k);
rebalance();
};
for(int i=1;i<=x+1;i++)add(nums[i]);
long long res=sum;
for(int i=x+2;i<nums.size();i++){
remove(nums[i-x-1]);
add(nums[i]);
res=min(res,sum);
}
return res;
}
后续待补充 中位数维护的两个例子，待我先做完 ✔️ 做完了
LeetCode480
固定长度为 k 的窗口，求滑动时每个区间的中位数，一般的思路，使用 big 维护前 \((k+1)/2\) 个元素
如果 k 为奇数，那么中位数为 *prev(big.end())
如果 k 为偶数，那么中位数为 (*small.begin()+(*prev(big.end())))/2
注意加和时溢出处理，其余都是固定套路
vector<double> cal(vector<int>& nums, int k) {
int n = nums.size();
multiset<int>big,small;
int m=(k+1)/2;
auto rebalance=[&](){
while(big.size()< m && !small.empty()){
auto it=small.begin();
big.insert(*it);
small.erase(it);
}
while(!big.empty() && !small.empty()){
auto r = prev(big.end()),l=small.begin();
int R=*r,L=*l;
if(R>L){
big.erase(r);
big.insert(L);
small.erase(l);
small.insert(R);
}else break;
}
};
auto erase_one=[&](int item){
if(auto it=big.find(item); it!=big.end()){
big.erase(it);
}else if(auto it=small.find(item); it!=small.end()){
small.erase(it);
}
};
auto add=[&](int k){
small.insert(k);
rebalance();
};
auto remove=[&](int k){
erase_one(k);
rebalance();
};
for(int i=0;i<k;i++)add(nums[i]);
vector<double>res;
double mid=(k%2?(*prev(big.end())):*small.begin())+1.0*(*prev(big.end()));
res.push_back(mid/2.0);
for(int i=k;i<n;i++){
remove(nums[i-k]);
add(nums[i]);
mid=(k%2?(*prev(big.end())):*small.begin())+1.0*(*prev(big.end()));
res.push_back(mid/2.0);
}
return res;
}
LeetCode295
这道题是窗口的长度是动态增大的，没有滑动过程，所以只需要维护add操作，不需要进行remove
class MedianFinder {
public:
multiset<int>big,small;
int k,m;
MedianFinder() {
k=0;
m=0;
}
void rebalance(){
while(big.size()< m && !small.empty()){
auto it=small.begin();
big.insert(*it);
small.erase(it);
}
while(!big.empty() && !small.empty()){
auto r = prev(big.end()),l=small.begin();
int R=*r,L=*l;
if(R>L){
big.erase(r);
big.insert(L);
small.erase(l);
small.insert(R);
}else break;
}
}
void addNum(int num) {
k++;
m+=k%2;
small.insert(num);
rebalance();
}
double findMedian() {
if(k%2==0){
return (*prev(big.end())+1.0*(*small.begin()))/2.0;
}else{
return *prev(big.end());
}
}
};
浙公网安备 33010602011771号