leetcode 3013 - 技术栈

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];
    }
};