LeetCode 805. 数组的均值分割

2024-05-10 HGNULB 约 11 分钟

文章信息 8 个分类 7 个标签

分类

题目描述

✅ 805. 数组的均值分割

题意分析

给一个整数数组 nums，问能否把它拆成两个都非空的子集 A 和 B（每个元素恰好归属其中一个），使得 A 的平均值等于 B 的平均值。只要返回能或不能，不用给出方案。

约束是解题的关键提示：数组长度不超过 30，元素值在 0 到 $10^4$ 之间。长度 30 意味着 $2^{30}$ 约十亿，纯枚举所有子集已经在超时边缘，但也说明「指数级但带剪枝」或者「按元素个数分层的可达性表」都还在预算内。元素非负且总和不超过 $3 \times 10^5$，说明「可能出现的子集和」这个集合是有界的，可以拿它当状态。

还有一层信号：题目问的是平均值相等，而平均值 = 和 ÷ 个数。这说明光记录「和」不够，必须同时记录「用了几个元素」，两个量绑在一起才能算出平均值。这一点直接决定了状态要开成二维。

边界：数组长度为 1 时无论如何拆不出两个非空子集，必须返回 false；元素允许为 0，[0, 0] 是合法的 true（两边平均值都是 0），任何「和为 0 就直接否定」的臆测剪枝都会在这里翻车；元素可以重复。

解法：子集和 DP

核心思路

暴力做法是枚举 A 的每一种取法，共 $2^n - 2$ 种（去掉全空和全取），对每种算出和与个数再比较两边平均值。$n = 30$ 时约十亿次，且每次还要累加，稳超时。

瓶颈在于枚举了「具体是哪些元素」，而判定只关心两个数字：A 的元素个数 $k$ 和 A 的元素和 $a$。大量不同的子集共享同一个 $(k, a)$，重复计算被白白浪费。

于是先做数学化简，把两个未知量压成一个。设总和为 sum、总个数为 n，A 有 $k$ 个元素、和为 $a$，那么 B 有 $n - k$ 个元素、和为 $sum - a$。两边平均值相等即

$\dfrac{a}{k} = \dfrac{sum - a}{n - k}$

交叉相乘得 $a(n - k) = k(sum - a)$，展开是 $an - ak = k \cdot sum - ka$，两边的 $ak$ 抵消，得到 $a \cdot n = k \cdot sum$，即

$a = \dfrac{sum \cdot k}{n}$

这个式子说明：两个子集平均值相等，等价于其中任意一个子集的平均值等于整个数组的平均值。判定条件因此从「二元关系」塌缩成「一元存在性」——只要存在某个 $k \in [1, n-1]$，使得 $sum \cdot k$ 能被 $n$ 整除，且能选出 $k$ 个元素其和恰为 $sum \cdot k / n$，答案就是 true。

剩下的问题就是「选恰好 $k$ 个元素能凑出哪些和」。定义状态：dp[k] = 从已经考虑过的元素中恰好选 $k$ 个，所能得到的全部和构成的集合。初值 dp[0] = {0}（一个都不选，和为 0），其余为空集，空集的含义是「这个个数目前不可达」。

转移是标准的 0-1 背包思路，只是背包多了一维「已选个数」：新来一个元素 num，对每个 $k$，dp[k] 应并入 {s + num | s ∈ dp[k-1]}——意思是「原先选了 $k-1$ 个凑出 $s$，现在把 num 也选上，就有了 $k$ 个凑出 $s + num$」。

循环不变量：处理完前 $i$ 个元素后，dp[k] 恰好等于「从这 $i$ 个元素中选 $k$ 个」所有可能和的集合，每个元素在同一个和里至多被用一次。要维持「至多用一次」，$k$ 必须从大到小更新：先算大的 $k$，它读取的 dp[k-1] 还是不含当前 num 的旧值；如果从小到大，dp[k-1] 已经吃进了 num，再传给 dp[k] 就等于把同一个 num 用了两次。这与一维 0-1 背包必须逆序遍历容量是同一个道理。

$k$ 的更新范围取 $[1, n-1]$ 而不是 $[1, n]$：dp[n] 代表整个数组都归 A，那样 B 就空了，题目不允许，索性不生成这个状态。

解题步骤

求 n 与 sum：后面的目标值公式 $sum \cdot k / n$ 全靠这两个量，先一次性算好。

建表并置初值：开 n + 1 个集合，只往 dp[0] 里放一个 0。这个 0 是所有转移的种子，漏了它整张表永远是空的。

外层遍历每个元素 num：这一层保证每个元素只有「选」或「不选」两种命运，是 0-1 而非完全背包。

内层 k 从 n-1 递减到 1：逆序是保证每个元素在同一个和里只用一次的唯一手段，写成正序就变成了可重复选取。

由 dp[k-1] 推 dp[k]：把 dp[k-1] 中每个和加上 num 存进 dp[k]。用集合而不是列表，天然去重——不同的元素组合凑出同一个和时只保留一份，这正是把指数级枚举压回多项式的地方。

逐个 k 检验：k 从 1 到 n-1，先判 sum * k % n != 0 就跳过（平均值不是整数和，凑不出来），再算 target = sum * k / n，若 dp[k] 含 target 立刻返回 true。整除判断必须在除法之前做，否则整数除法会把 target 悄悄取整成一个错误的值。

全部落空返回 false。

以 nums = [1, 2, 3, 4] 走一遍。n = 4，sum = 10，平均值 2.5。初始 dp[0] = {0}，dp[1] = dp[2] = dp[3] = {}。

处理 num = 1：k = 3 读 dp[2] 为空，无事发生；k = 2 读 dp[1] 为空，无事发生；k = 1 读 dp[0] = {0}，得 dp[1] = {1}。
处理 num = 2：k = 3 读 dp[2] 仍为空；k = 2 读 dp[1] = {1}，得 dp[2] = {3}；k = 1 读 dp[0] = {0}，dp[1] 变为 {1, 2}。
处理 num = 3：k = 3 读 dp[2] = {3}，得 dp[3] = {6}；k = 2 读 dp[1] = {1, 2}，dp[2] 变为 {3, 4, 5}；k = 1 得 dp[1] = {1, 2, 3}。
处理 num = 4：k = 3 读 dp[2] = {3, 4, 5}，dp[3] 变为 {6, 7, 8, 9}；k = 2 读 dp[1] = {1, 2, 3}，dp[2] 变为 {3, 4, 5, 6, 7}；k = 1 得 dp[1] = {1, 2, 3, 4}。

检验阶段：k = 1 时 10 × 1 % 4 = 2 ≠ 0，跳过；k = 2 时 20 % 4 = 0，target = 5，而 dp[2] = {3, 4, 5, 6, 7} 含 5，返回 true。对应的方案是 {1, 4} 与 {2, 3}，平均值都是 2.5。

顺带注意 k = 2 那一轮的逆序意义：处理 num = 4 时先做 k = 3 再做 k = 2 最后做 k = 1。若顺序反过来，dp[1] 会先被写入 4，紧接着 k = 2 读到这个 4 又加一次 4，dp[2] 里会冒出 8——而真实的两元素最大和只有 7，这个 8 是「同一个 4 用了两次」的伪状态。

代码实现

class Solution {
    // 用 dp[k] 记录选 k 个数能得到的所有和。
    public boolean splitArraySameAverage(int[] nums) {
        int n = nums.length;
        int sum = 0;
        for (int v : nums) {
            sum += v;
        }

        List<Set<Integer>> dp = new ArrayList<>();
        for (int i = 0; i <= n; i++) {
            dp.add(new HashSet<>());
        }
        dp.get(0).add(0);

        for (int num : nums) {
            for (int k = n - 1; k >= 1; k--) {
                for (int s : dp.get(k - 1)) {
                    dp.get(k).add(s + num);
                }
            }
        }

        for (int k = 1; k < n; k++) {
            if (sum * k % n != 0) {
                continue;
            }
            int target = sum * k / n;
            if (dp.get(k).contains(target)) {
                return true;
            }
        }

        return false;
    }
}

func splitArraySameAverage(nums []int) bool {
    // 用 dp[k] 记录选 k 个数能得到的所有和。
    n := len(nums)
    sum := 0
    for _, v := range nums {
        sum += v
    }

    dp := make([]map[int]struct{}, n+1)
    dp[0] = make(map[int]struct{})
    dp[0][0] = struct{}{}

    for _, num := range nums {
        for k := n - 1; k >= 1; k-- {
            if dp[k] == nil {
                dp[k] = make(map[int]struct{})
            }
            for s := range dp[k-1] {
                dp[k][s+num] = struct{}{}
            }
        }
    }

    for k := 1; k < n; k++ {
        if sum*k%n != 0 {
            continue
        }
        target := sum * k / n
        if _, ok := dp[k][target]; ok {
            return true
        }
    }

    return false
}

复杂度分析

时间复杂度：$O(n^2 \cdot S)$，其中 $n$ 是数组长度、$S$ 是单层 dp[k] 中可达和的数量上界。外层遍历 $n$ 个元素，内层遍历 $n$ 个个数维，最里层遍历一个集合，集合大小被「可达和的取值范围」而非「子集个数」限制——这正是用集合去重换来的收益，否则最里层就是 $C(n, k)$ 级别的指数量。$S$ 同时受 $\min(2^{n}, sum)$ 约束，本题 $n \le 30$、$sum \le 3 \times 10^5$，实测最坏输入在两百毫秒内跑完。

空间复杂度：$O(n \cdot S)$，$n + 1$ 个集合，每个最多装 $S$ 个和。凭的是状态必须同时记录「个数」和「和」两维，任何一维都不能省——省掉个数维就无法把和与目标 $sum \cdot k / n$ 中的 $k$ 对上号。

关键点总结

「两个子集平均值相等」通过一次交叉相乘化简为「某个子集的平均值等于全局平均值」，把二元条件压成一元存在性判定。这类均值 / 比例题第一步永远是先做代数化简，别急着写搜索。

平均值由「和」与「个数」两个量共同决定，所以状态必须是二维的 (个数, 和)。凡是判定条件里出现除法的题，分子分母都要进状态。

0-1 背包的「每个物品只用一次」靠逆序更新保证，这是可以直接迁移的模板级结论：正序等于完全背包。

用集合存「可达和」而不是列表存「具体方案」，是把 $C(n, k)$ 级的枚举压成 $S$ 级的关键；去重发生的地方就是复杂度下降的地方。

整除性判断 sum * k % n == 0 是一个极强的剪枝，很多 k 根本不用查表。写在除法之前既是剪枝也是正确性保证。

面试视角：这题的分水岭在于能不能当场推出 $a = sum \cdot k / n$。推不出来就只能写指数搜索，推出来之后剩下的部分是标准背包。面试时建议先把这条公式写在白板上讲清楚再动手，然后主动补充两个优化方向——一是只需枚举 $k \le n / 2$（A 和 B 里必有一个大小不超过一半，答案对称），二是当 $n \le 30$ 时可以改用折半枚举（meet in the middle）把 $2^{30}$ 降到 $2^{15}$ 级别。能说出「先化简再背包」的思路脉络，比背下代码有用得多。

易错点总结

内层 k 写成正序递增：处理 [1, 2, 3, 4] 的元素 4 时，dp[1] 先被写入 4，紧接着 k = 2 又把这个 4 加一次，dp[2] 里出现 8，而两个元素的真实最大和只有 7；一旦某个 target 命中这类伪和就会错误返回 true。

检验循环写成 k <= n：[1, 2] 中 k = 2 时 3 × 2 % 2 = 0、target = 3，若 dp[2] 被生成且含 3 就返回 true，但那意味着 B 是空集，正确答案是 false。

检验循环从 k = 0 开始：sum × 0 % n 恒为 0，target = 0，而 dp[0] 永远含 0，于是任何输入都返回 true，[1, 2] 直接答错。

跳过整除判断直接整数除法：[1, 2] 中 k = 1 时 3 / 2 被截断成 1，而 dp[1] = {1, 2} 含 1，返回 true；但 [1] 和 [2] 的平均值分别是 1 和 2，正确答案是 false。

丢掉「个数」这一维，退化成普通子集和判定：[2, 3, 4, 9] 的正确答案是 false（k 必须为偶数，k = 2 的 target = 9，两元素和只有 5、6、7、11、12、13）；但如果只问「存在子集和为 9 吗」，单个元素 9 就命中了，错误返回 true。

忘记 dp[0].add(0)：所有集合都是空的，转移无米下锅，[1, 2, 3, 4] 这种明明可行的输入也返回 false。

转移时忘记加 num（写成 dp[k].addAll(dp[k-1])）：[1, 2, 3, 4] 的 dp[2] 会退化成 dp[1] 的副本 {1, 2, 3, 4}，target = 5 找不到，返回 false。

Go 里忘记给 dp[k] 惰性初始化：dp 由 make 出来时每个槽位都是 nil map，直接写 dp[k][s+num] = struct{}{} 会 panic: assignment to entry in nil map；读 nil map 是安全的，只有写才崩，所以这个坑只在写入侧出现。

加上「sum == 0 直接返回 false」的臆想剪枝：[0, 0] 的两个子集平均值都是 0，正确答案是 true，这条剪枝会答错。题目允许元素为 0。

长度为 1 时额外写特判但写反：n = 1 时检验循环 k 从 1 到 n-1 = 0 本就不会执行，天然返回 false；自己加的特判反而容易写成返回 true。

相似题目

题目	难度	考察点
416. 分割等和子集	中等	目标和固定为 `sum / 2` 且不限制元素个数，状态只需一维，是本题去掉个数维后的形态
494. 目标和	中等	每个元素必选但可正可负，先化简成「选一个子集凑出定值」再计数，求的是方案数而非可行性
698. 划分为k个相等的子集	中等	要拆成 k 份而不是 2 份，子集和 DP 不够用，需要状态压缩加回溯
473. 火柴拼正方形	中等	698 中 k 固定为 4 的特例，考点转向剪枝顺序与去重
1049. 最后一块石头的重量 II	中等	同样先做代数化简得到「两堆差值最小」，但求的是最优值而非存在性
LCR 101. 分割等和子集	简单	与 416 同题，可直接套用一维背包模板