3013: 将数组分成最小总代价的子数组Ⅱ
第一段的第一个数是确定的,即 nums[0]。
如果知道了第二段的第一个数的位置(记作 p),第三段的第一个数的位置,......,第 k 段的第一个数的位置(记作 q),那么这个划分方案也就确定了。
考虑到 q−p≤dist,本题相当于++在一个大小固定为 dist+1 的滑动窗口内,求前 k−1 小的元素和++。
用两个有序集合来做:
- 初始化两个有序集合 L 和 R。注意:为方便计算,把 k 减 1。
- 把 nums[1] 到 nums[dist+1] 加到 L 中。
- 保留 L 最小的 k 个数,把其余数丢到 R 中。
- 从 i=dist+2 开始滑窗。
- 先把 out=nums[i−dist−1] 移出窗口:如果 out 在 L 中,就从 L 中移除,否则从 R 中移除。
- 然后把 in=nums[i] 移入窗口:++如果 in 小于 L 中的最大元素++,则加入 L,否则加入 R。
- 上面两步做完后,如果 L 中的元素个数小于 k(等于 k−1),则从 R 中取一个最小元素加入 L;反之,如果 L 中的元素个数大于 k(等于 k+1),则从 L 中取一个最大元素加入 R。
上述过程维护 L 中元素之和 sumL,取 sumL 的最小值,即为答案。
sort(nums.begin()+1,nums.end()); //排序
long long sum=reduce(nums.begin()+1,nums.begin()+dist+2,0LL); //求和
multiset<int> L(nums.begin()+1,nums.begin()+dist+2)
multiset<int> :默认升序,允许重复值。
-
nums.begin()+1:起始迭代器(跳过第 0 个元素,从索引 1 开始) -
nums.begin()+dist+2:结束迭代器(不包含 该位置,实际累加到索引dist+1)int x=*L.rbegin();
int x=*R.begin();

class Solution {
public:
long long minimumCost(vector<int>& nums, int k, int dist) {
k--;
//sumL,nums[1]到nums[dist+1]之和
long long sum=reduce(nums.begin()+1,nums.begin()+dist+2,0LL);
multiset<int> L(nums.begin()+1,nums.begin()+dist+2),R;
auto L2R=[&](){
int x=*L.rbegin();
sum-=x;
L.erase(L.find(x));
R.insert(x);
};
auto R2L=[&](){
int x=*R.begin();
sum+=x;
R.erase(R.find(x));
L.insert(x);
};
while(L.size()>k) L2R(); //保留 L 最小的 k 个数(注意开头已经把 k 减 1)
long long ans=sum;
for(int i=dist+2;i<nums.size();i++){
//移除 out
int out=nums[i-dist-1];
auto it=L.find(out);
if(it!=L.end()){ //如果 out 在 L 中
sum-=out;
L.erase(it);
}
else R.erase(R.find(out));
//添加 in
int in=nums[i];
if(in<*L.rbegin()){ //如果 in 小于 L 中的最大元素
sum+=in;
L.insert(in);
}
else R.insert(in);
//维护大小
if(L.size()<k) R2L();
else if(L.size()>k) L2R();
ans=min(ans,sum); //维护 L 中元素之和 sumL,取 sumL 的最小值
}
return ans+nums[0];
}
};