【LeetCode 每日一题】3013. 将数组分成最小总代价的子数组 II

Problem: 3013. 将数组分成最小总代价的子数组 II

文章目录

[1. 整体思路](#1. 整体思路)
- - 核心问题
  - [算法逻辑：滑动窗口 + 双堆/平衡树维护 Top K](#算法逻辑：滑动窗口 + 双堆/平衡树维护 Top K)
[2. 完整代码](#2. 完整代码)
[3. 时空复杂度](#3. 时空复杂度)
- - [时间复杂度： O ( N log ⁡ ( dist ) ) O(N \log (\text{dist})) O(Nlog(dist))](#时间复杂度： O ( N log ⁡ ( dist ) ) O(N \log (\text{dist})) O(Nlog(dist)))
  - [空间复杂度： O ( dist ) O(\text{dist}) O(dist)](#空间复杂度： O ( dist ) O(\text{dist}) O(dist))

1. 整体思路

核心问题

我们需要将数组 nums 分成 k 个子数组。

总代价 = 第 1 个子数组的头 + 第 2 个子数组的头 + ... + 第 k 个子数组的头。

限制条件：

第 1 个子数组必须从 nums[0] 开始，所以 nums[0] 必选。
第 2 到第 k 个子数组的头，必须在 nums 中选取。
关键限制 dist ：这些被选中的头元素之间的索引差不能超过 dist。具体来说，如果我们选了索引 i 1 , i 2 , ... , i k − 1 i_1, i_2, \dots, i_{k-1} i1,i2,...,ik−1 作为剩下 k − 1 k-1 k−1 个子数组的头，那么必须满足 i j + 1 − i j ≤ d i s t i_{j+1} - i_j \le dist ij+1−ij≤dist。
- 修正：实际上题目的意思是，除了 nums[0] 外，我们需要在 nums 的剩余部分中，选取 k − 1 k-1 k−1 个元素，使得这 k − 1 k-1 k−1 个元素在原数组中的索引跨度不超过 dist（或者更准确地说，是一个长度为 dist+1 的滑动窗口内选 k − 1 k-1 k−1 个最小的数）。
- 题目转化为：在 nums[1] 到 nums[n-1] 中，找到一个长度为 dist + 1 的滑动窗口，在这个窗口内选出最小的 k-1 个数，使得它们的和最小。

算法逻辑：滑动窗口 + 双堆/平衡树维护 Top K

这个问题等价于 "滑动窗口中的最小 k 个数之和"。

数据结构设计：
- 我们需要动态维护一个滑动窗口内的数值，并能够快速知道其中最小的 k-1 个数的和。
- 使用两个 TreeMap（平衡二叉搜索树）来模拟两个集合：
  - left (Top K set) ：存储窗口中最小的 k-1 个数。
  - right (Candidate set) ：存储窗口中剩余的数 （都比 left 中的大）。
- 维护变量 sum：记录 left 中所有元素的和。
- 维护变量 leftSize：记录 left 中的元素个数。
滑动窗口流程：
- 初始化 ：先将前 dist + 1 个元素（从 nums[1] 开始）加入数据结构，并调整 left 和 right 使得 left 恰好包含最小的 k-1 个数。
- 滑动：
  1. 移除元素 out ：窗口左移，移除最左边的元素（nums[i - dist - 1]）。
    - 如果在 left 中，更新 sum 和 leftSize。
    - 如果在 right 中，直接移除。
  2. 添加元素 in ：窗口右移，加入新元素 nums[i]。
    - 如果 in 比 left 中的最大值小（有资格进入 Top K），加入 left，更新 sum。
    - 否则加入 right。
  3. 再平衡 (Rebalance) ：
    - 如果 leftSize < k-1：从 right 中取最小值补充到 left。
    - 如果 leftSize > k-1：将 left 中的最大值移到 right。
  4. 更新答案 ：ans = min(ans, sum)。

2. 完整代码

java 复制代码

import java.util.TreeMap;

class Solution {
    // left 存储滑动窗口中最小的 k-1 个数
    // key: 数值, value: 出现次数 (处理重复元素)
    private final TreeMap<Integer, Integer> left = new TreeMap<>(); 
    
    // right 存储滑动窗口中剩余的数
    private final TreeMap<Integer, Integer> right = new TreeMap<>();
    
    // sum 记录 left 中所有元素的和
    private long sum;    
    
    // leftSize 记录 left 中元素的总个数 (考虑了 value 中的计数)
    private int leftSize;  

    public long minimumCost(int[] nums, int k, int dist) {
        // 因为 nums[0] 必选，我们还需要选 k-1 个数
        // 所以将 k 减 1，后续逻辑中的 k 即代表需要在滑动窗口中选出的数的个数
        k -= 1; 

        // 初始化 sum 为 nums[0] (必选代价)
        sum = nums[0];
        left.clear();
        right.clear();
        leftSize = 0;

        // 1. 初始化窗口：处理前 dist + 1 个元素 (nums[1] ... nums[dist+1])
        // 先把所有元素一股脑加进 left
        for (int i = 1; i <= dist + 1; i++) {
            add(left, nums[i]);
            leftSize++;
            sum += nums[i];
        }

        // 如果 left 中的元素超过了 k 个，把最大的移到 right，直到只剩 k 个
        // 这样初始化后，left 里就是窗口内最小的 k 个数
        while (leftSize > k) moveLeftToRight();

        // 记录当前最小代价
        long ans = sum;

        // 2. 开始滑动窗口
        // i 代表新进入窗口的元素索引
        for (int i = dist + 2; i < nums.length; i++) {
            // out 是要滑出窗口的元素
            int out = nums[i - dist - 1];

            // 移除 out
            if (contains(left, out)) {
                remove(left, out);
                leftSize--;
                sum -= out;
            } else {
                remove(right, out);
            }

            // in 是新进入窗口的元素
            int in = nums[i];
            
            // 决定把 in 加到哪里
            if (leftSize == 0) {
                // 如果 left 为空 (k=0 的情况不应该发生，但作为防御性编程)，直接加 right
                add(right, in);
            } else {
                // 如果 in 比 left 中最大的数还小，说明它有资格进入 Top K
                int leftMax = left.lastKey();
                if (in < leftMax) {
                    add(left, in);
                    leftSize++;
                    sum += in;
                } else {
                    add(right, in);
                }
            }

            // 再平衡：确保 left 中恰好有 k 个数
            // 情况 1: left 不够 k 个 (可能是刚才移除了一个 left 中的元素，而新加的进了 right)
            while (leftSize < k) moveRightToLeft();
            
            // 情况 2: left 超过 k 个 (可能是刚才移除了 right 中的元素，而新加的进了 left)
            while (leftSize > k) moveLeftToRight();

            // 更新全局最小代价
            ans = Math.min(ans, sum);
        }

        return ans;
    }

    // 将 left 中最大的元素移动到 right
    private void moveLeftToRight() {
        int x = left.lastKey();
        remove(left, x);
        leftSize--;
        sum -= x;
        add(right, x);
    }

    // 将 right 中最小的元素移动到 left
    private void moveRightToLeft() {
        int x = right.firstKey();
        remove(right, x);
        add(left, x);
        leftSize++;
        sum += x;
    }

    // 辅助：向 TreeMap 添加元素
    private void add(TreeMap<Integer, Integer> mp, int x) {
        mp.merge(x, 1, Integer::sum);
    }

    // 辅助：从 TreeMap 移除元素
    private void remove(TreeMap<Integer, Integer> mp, int x) {
        int c = mp.get(x);
        if (c == 1) mp.remove(x);
        else mp.put(x, c - 1);
    }

    // 辅助：检查是否存在
    private boolean contains(TreeMap<Integer, Integer> mp, int x) {
        return mp.containsKey(x);
    }
}

3. 时空复杂度

假设数组长度为 N N N。

时间复杂度： O ( N log ⁡ ( dist ) ) O(N \log (\text{dist})) O(Nlog(dist))

计算依据 ：
- 代码主体是一个循环，遍历数组一次。
- 在循环内部，涉及 TreeMap 的操作（add, remove, lastKey, firstKey）。
- TreeMap 中的元素数量最多为滑动窗口的大小，即 dist + 1。
- 因此，单次操作的复杂度为 O ( log ⁡ ( dist ) ) O(\log (\text{dist})) O(log(dist))。
结论： O ( N log ⁡ ( dist ) ) O(N \log (\text{dist})) O(Nlog(dist))。

空间复杂度： O ( dist ) O(\text{dist}) O(dist)

计算依据 ：
- 两个 TreeMap (left, right) 存储的元素总数就是滑动窗口的大小，即 dist + 1。
结论： O ( dist ) O(\text{dist}) O(dist)。