3013.将数组分成最小总代价的子数组 II：两个堆维护k-1小 + 滑动窗口

【LetMeFly】3013.将数组分成最小总代价的子数组 II：两个堆维护k-1小 + 滑动窗口

力扣题目链接：https://leetcode.cn/problems/divide-an-array-into-subarrays-with-minimum-cost-ii/

给你一个下标从 0 开始长度为 n 的整数数组 nums 和两个 正 整数 k 和 dist 。

一个数组的 代价 是数组中的 第一个 元素。比方说，[1,2,3] 的代价为 1 ，[3,4,1] 的代价为 3 。

你需要将 nums 分割成 k 个 连续且互不相交 的子数组，满足 第二 个子数组与第 k 个子数组中第一个元素的下标距离 不超过 dist 。换句话说，如果你将 nums 分割成子数组 nums[0..(i1 - 1)], nums[i1..(i2 - 1)], ..., nums[ik-1..(n - 1)] ，那么它需要满足 ik-1 - i1 <= dist 。

请你返回这些子数组的 最小 总代价。

 

示例 1：

输入：nums = [1,3,2,6,4,2], k = 3, dist = 3
输出：5
解释：将数组分割成 3 个子数组的最优方案是：[1,3] ，[2,6,4] 和 [2] 。这是一个合法分割，因为 ik-1 - i1 等于 5 - 2 = 3 ，等于 dist 。总代价为 nums[0] + nums[2] + nums[5] ，也就是 1 + 2 + 2 = 5 。
5 是分割成 3 个子数组的最小总代价。

示例 2：

输入：nums = [10,1,2,2,2,1], k = 4, dist = 3
输出：15
解释：将数组分割成 4 个子数组的最优方案是：[10] ，[1] ，[2] 和 [2,2,1] 。这是一个合法分割，因为 ik-1 - i1 等于 3 - 1 = 2 ，小于 dist 。总代价为 nums[0] + nums[1] + nums[2] + nums[3] ，也就是 10 + 1 + 2 + 2 = 15 。
分割 [10] ，[1] ，[2,2,2] 和 [1] 不是一个合法分割，因为 ik-1 和 i1 的差为 5 - 1 = 4 ，大于 dist 。
15 是分割成 4 个子数组的最小总代价。

示例 3：

输入：nums = [10,8,18,9], k = 3, dist = 1
输出：36
解释：将数组分割成 4 个子数组的最优方案是：[10] ，[8] 和 [18,9] 。这是一个合法分割，因为 ik-1 - i1 等于 2 - 1 = 1 ，等于 dist 。总代价为 nums[0] + nums[1] + nums[2] ，也就是 10 + 8 + 18 = 36 。
分割 [10] ，[8,18] 和 [9] 不是一个合法分割，因为 ik-1 和 i1 的差为 3 - 1 = 2 ，大于 dist 。
36 是分割成 3 个子数组的最小总代价。

 

提示：

• 3 <= n <= 105
• 1 <= nums[i] <= 109
• 3 <= k <= n
• k - 2 <= dist <= n - 2

解题方法：有序集合 + 滑动窗口

$nums$第一个元素必选，剩下$k-1$个元素的起始位置间隔必须$\leq dist$。使用一个大小为$dist+1$的滑动窗口，每次求这个窗口中$k-1$小元素的和。

问题变成了滑动窗口向右滑动过程中，窗口中不断新增一个元素、移除一个元素的状态下如何保持计算$k-1$小的元素有哪些：

我们可以使用两个有序集合，一个叫$stage$代表(正在舞台上的)$k-1$小元素，一个叫$candidate$代表在窗口中但(暂)未被选中的元素。

窗口右移过程中，假设要新加入窗口的元素是$in$，移除窗口的元素是$out$，

对于$out$有：

• 如果$out$在$candidate$候选集合中，那么$out$永无上台之日，直接移出候选
• 如果$out$在$stage$选中集合中，那么$out$是时候退役了，移出$stage$集合，并更新集合中元素之和

对于$in$有：

• 如果$in$比$stage$中最大元素小，说明更优秀的人来了，移出$stage$集合中最大的那个，将$in$加入$stage$集合，并更新$stage$集合中元素之和
• 否则，将$in$加入候选

之后进行$stage$集合大小的调整：

• 如果$stage$集合小于$k-1$，说明有人退役暂无人顶上，从$candidate$中选最小的那个顶上(移出$candidate$并加入$stage$)，并更新$stage$集合中元素之和
• 如果$stage$集合大于$k+1$，说明有更优秀的人来了，要把$stage$中最大的那个移出并加入到$candidate$，并更新$stage$集合中元素之和

整个滑动窗口移动的过程中，所有$stage$元素和中最小的那个即为答案。

为了方便计算，我们开局直接把$k$减一。

• 时间复杂度$O(n\log dist)$
• 空间复杂度$O(dist)$

AC代码

C++

/*
 * @LastEditTime: 2026-02-03 22:11:11
 */
typedef long long ll;
class Solution {
public:
    ll minimumCost(vector<int>& nums, int k, int dist) {
        k--;
        multiset<ll> stage(nums.begin() + 1, nums.begin() + dist + 2), candidate;
        ll ans = accumulate(nums.begin(), nums.begin() + dist + 2, 0ll);
        while (stage.size() > k) {
                int retiree = *stage.rbegin();
                stage.erase(prev(stage.end()));
                ans -= retiree;
                candidate.insert(retiree);
        }

        ll nowAns = ans;
        for (int end = dist + 2; end < nums.size(); end++) {
            int in = nums[end], out = nums[end - dist - 1];
            
            // out
            multiset<ll>::iterator it = candidate.find(out);
            if (it != candidate.end()) {
                candidate.erase(it);
            } else {
                stage.erase(stage.find(out));
                nowAns -= out;
            }

            // in
            if (in < *stage.rbegin()) {
                stage.insert(in);
                nowAns += in;
            } else {
                candidate.insert(in);
            }

            // resize
            if (stage.size() == k - 1) {
                int newActor = *candidate.begin();
                candidate.erase(candidate.begin());
                stage.insert(newActor);
                nowAns += newActor;
            } else if (stage.size() == k + 1) {
                int retiree = *stage.rbegin();
                stage.erase(prev(stage.end()));
                nowAns -= retiree;
                candidate.insert(retiree);
            }

            ans = min(ans, nowAns);
        }
        return ans;
    }
};

#if defined(_WIN32) || defined(__APPLE__)
/*
[1,3,2,6,4,2]
3
3

5
*/
int main() {
    string s;
    int a, b;
    while (cin >> s >> a >> b) {
        Solution sol;
        vector<int> v = stringToVector(s);
        cout << sol.minimumCost(v, a, b) << endl;
    }
    return 0;
}
#endif

同步发文于CSDN和我的个人博客，原创不易，转载经作者同意后请附上原文链接哦~

千篇源码题解已开源