每日一题 —— LC. 805 数组的均值分割

抠脚的大灰狼

已于 2022-11-17 08:36:59 修改

阅读量537

点赞数

分类专栏：算法文章标签：算法数据结构二进制压缩位运算

于 2022-11-16 21:16:14 首次发布

本文链接：https://blog.csdn.net/vcj1009784814/article/details/127893559

版权

算法专栏收录该内容

120 篇文章

订阅专栏

该博客探讨了一个数组均值分割问题，即如何将数组分成两部分，使得两部分的平均值相等。博主介绍了三种解题思路：1）使用折半搜索；2）通过变形处理数组后再进行折半搜索；3）利用动态规划和二进制优化。每种方法都详细解释了思路和代码实现，并分析了时间复杂度。最后，博主指出可以进一步优化动态规划的实现，减少时间复杂度。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

805. 数组的均值分割

给定你一个整数数组 nums

我们要将 nums 数组中的每个元素移动到 A 数组或者 B 数组中，使得 A 数组和 B 数组不为空，并且 average(A) == average(B) 。

如果可以完成则返回true ，否则返回 false 。

注意：对于数组 arr , average(arr) 是 arr 的所有元素的和除以 arr 长度。

提示:

1 <= nums.length <= 30
0 <= nums[i] <= 10^4

思路

题目要求将整个数组分成两部分，两部分的平均值相等。我们假设能够完成这样的划分，设A组和B组的平均值为avg，那么容易得知，整个数组的平均值也是avg。（我们可以把A组和B组的每个元素都替换成avg，则整个数组的总和不变，那么显而易见，每个元素都是avg，那整个数组的平均值就是avg）。

这是第一个比较关键的信息，即若存在这样的划分，则A组的平均数=B组的平均数=整个数组的平均数。

设A组的总和为sumA，元素个数为k，设整个数组的总和为sum，元素个数为n。

若能完成这样的划分，则必须满足：sumA / k = (sum - sumA) / (n - k) = sum / n

这个式子我们可以只取一半，即 sumA / k = sum / n，只要满足这个式子，那么剩余的数组成的B组，其平均值一定是avg。

简单证明一下，如果存在某个A组满足 sumA / k = sum / n，则B组的总和是 sum - sumA = sum - k * avg，B组的元素个数为n - k，B组的平均值是 (sum - k * avg) / (n - k)，由于sum = avg * n，我们进行替换后得 (n * avg - k * avg) / (n - k) = avg。

所以，我们对题目进行一下提炼，问题就变成了，能否从整个数组中取出某些数，组成A数组，设这部分数的总和为sumA，数的个数为k，使得其满足 sumA / k = sum / n。

思路一

折半搜索

既然是选取某些数，那么容易想到，每个数都有选或者不选两种情况，根据题目的数据范围，数组的总长度最大为30，那么总的方案数就是 $2^{30}$ ，如果直接暴力枚举所有方案，则时间复杂度大于在 $10^9$ 级别，是一定会超时的。

这时我们可以使用一种技巧，叫做折半查找。假设数的总个数为30，我们可以先对前15个数做一次暴力搜索，找出前15个数的所有组合方案，这时的复杂度为 $2^{15} ≈ 3×10^4$ ，将搜索结果保存在一个哈希表当中；然后对后15个数再做一次暴力搜索，对于后15个数的每种方案，我们看一下能否和前15个数的某个方案组合在一起，形成满足条件的方案。这样，总的复杂度就从 $2^{30}$ 缩减为了 $2 × 2^{15}$ ，大概在 $10^5$ 级别以内，就不会超时啦。

接下来看一下具体操作。

我们在对前一半的元素做搜索时，保存下每种方案的总和 $t o t$ ，以及元素的个数 $c n t$ ，即，将每种方案保存为一个二元组。（因为我们需要判断的条件是 sumA / k = sum / n，需要和以及个数）

然后，再对后一半的元素做搜索，假设后一半元素的某个选择方案，得到的和是 $tot^{'}$ ，元素个数为 $cnt^{'}$ ，

假设前一半元素的方案中，存在某个方案，能够与当前这个方案组合起来，成为满足条件的一种方案。则前半元素中的这个方案，假设其和为 $t o t$ ，个数为 $c n t$ 。那么一定有如下的等式： $\frac{tot + tot^{'}}{cnt + cnt^{'}} = \frac{sumA}{k} = \frac{sum}{n}$

那么我们如何找到，在前半个元素的方案中，是否存在这样一种方案呢？答案是枚举k，k的有效范围仅仅是[1, n - 1]（枚举k的最大循环次数也就是30），那么对于后半元素的当前方案（和为 $t o t$ ，个数为 $c n t$ ），我们只需要枚举一下k，每次计算 $cnt^{'} = k - cnt$ ，看一下哈希表中，是否存在个数为 $cnt^{'}$ 的方案，若存在个数为 $cnt^{'}$ 的方案，看一下这些方案中，是否有某个方案的和等于 $tot^{'}$ ，（ $tot^{'} = \frac{sum × k}{n} - tot$ ），若有，则找到一种有效划分。若遍历了所有的情况，都米有找到，则不存在有效的划分。

那么我们的哈希表，需要以个数 $c n t$ 作为 key，value则是所有个数为 $c n t$ 的方案的和。

这样总的时间复杂度，我们稍微算一下，前半元素的查找需要 $2^{15} ≈ 3 × 10^4$ ；后半元素的查找需要 $2^{15}$ ，每种方案需要枚举一下k，来查找是否在前半部分中存在与之匹配的方案，则后半部分的查找总复杂度是 $2^{15} × 30 ≈ 10^6$ ，总的复杂度大概在 $10^6$ 级别，那么不会超时。

// C++ 124ms
class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        int n = nums.size();
        if (n <= 1) return false; // 长度为1, 则一定无法划分
        int sum = 0;
        for (auto& i : nums) sum += i;
        // cnt作为key, 因为cnt最大为30, 这里直接开到40, 用数组来模拟哈希表
        unordered_set<int> f[40];
        int m = n / 2; // 总的个数的一半
        // 枚举前一半的数, 所有选择的方案状态 (用二进制表示)
        // 比如 1011 表示选择了 nums[0], nums[1], nums[3]
        for (int st = 0; st < (1 << m); st++) {
            int tot = 0, cnt = 0;
            // 对于当前的st 这个状态, 看一下[0, m - 1], 共m个位置上, 那些数被选择了
            for (int j = 0; j < m; j++) {
                // 若nums[j]被选择了, 则计数并累加
                if ((st >> j) & 1) tot += nums[j], cnt++;
            }
            // cnt为key, 向 set 中插入一个 tot
            f[cnt].emplace(tot);
        }

        // 对余下的一半的数, 个数为 n - m, 进行搜索
        for (int st = 0; st < (1 << (n - m)); st++) {
            int tot = 0, cnt = 0;
            for (int j = 0; j < n - m; j++) {
                if ((st >> j) & 1) tot += nums[j + m], cnt++;
            }
            // k = tot + tot'
            // 由于A和B不能有一个为空, 则k的有效取值范围是 [1, n - 1]
            // sumA / k = sum / n
            // sumA = sum * k / n 
            for (int k = max(1, cnt); k < n; k++) {
                if (sum * k % n) continue; // sumA一定是个整数, 所以当计算出来的sumA不是整数时，直接跳过
                if (f[k - cnt].count(sum * k / n - tot)) return true;
            }
        }
        return false;
    }
};

思路二

变形处理 + 折半搜索

接思路一，我们可以这样优化：将整个数组中的每个元素，都减去一个avg，这样以来，整个数组的平均值就变成了0。

而由于avg = sum / n，可能是浮点数，为了不产生浮点数，我们对每个数都乘上一个n，这样整个数组的avg就等于原先的sum。

让整个数组的平均值变成0有什么好处呢？当然是方便判断啦。

剩下的部分，和思路一一样，采用折半查找。不过我们的判断条件就变得更为简单了，假设存在满足条件的划分，那么我们找到的A数组的和一定为0，B数组的和也一定为0。

所以，在我们对原数组进行了变形处理后，问题就变成了：能否从数组中选取一部分数出来，使得这部分的数的和为0。

同样采用折半搜索，当我们处理后半部分的数时，对于某种方案，其和为tot'，只需要在前半部分的方案中，找到一个和为-tot的，即可完成划分。并且，在搜索的过程中，如果遇到了某种方案的tot = 0，那么可以直接提前返回。另外，需要注意，最终的方案不能选取数组中全部的数（这样会使得另一个B数组为空，不满足题目要求）。

// C++ 152ms
class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        int n = nums.size();
        if (n <= 1) return false;
        int sum = accumulate(nums.begin(), nums.end(), 0); // 计算和
        // 每个数都减去平均值
        for (auto& i : nums) i = i * n - sum;
        unordered_set<int> left;
        int m = n / 2;
        // s 从1开始, 不从0开始, 0表示一个都不选, 会增加后续代码的判断逻辑的复杂性
        for (int s = 1; s < (1 << m); s++) {
            int tot = 0;
            for (int j = 0; j < m; j++) {
                if ((s >> j) & 1) tot += nums[j];
            }
            if (tot == 0) return true; // s != 0, 不能一个也不选
            left.emplace(tot); // 同样是往set里插入元素, 据说 emplace 比 insert 效率高
        }

        // 后半部分全选时的和
        int right_sum = 0;
        for (int i = m; i < n; i++) right_sum += nums[i];

        // 后半部分也不从0开始, 若后半部分为0, 即一个不选时, 也能满足条件, 则说明前半部分全选时tot=0, 在之前就已经提前退出了
        for (int s = 1; s < (1 << (n - m)); s++) {
            int tot = 0;
            for (int j = 0; j < n - m; j++) {
                if ((s >> j) & 1) tot += nums[j + m];
            }
            if (tot == 0) return true;
            //if (s.count(-tot)) return true; // 这样写不对, 可能是右半部分全选了, 左半部分有全选了
            // 怎样判断左半部分和右半部分都全选了呢?
            // 容易知道, 如果右半部分全选, 则此时的右半部分的和为tot, 且在set中一定存在一个 -tot (左半部分全选) , 因为整个数组的和为0
            // 此时 tot == right_sum
            // 那么 tot == right_sum 是否一定是在右半部分的数全选时才能达到呢? 
            // 换句话说, 有没有可能右半部分的数没有全选, 就满足了 tot == right_sum 呢？
            // 这是有可能的, 但是！此时右半部分剩余没选的那些数, 它们的和一定是0, 因为选不选这些数, 总和不变
            // 而这部分的数的和为0, 一定出现在 tot == right_sum 之前, 也就是说, 会提前就遇到这个情况, 就返回了
            // 因为右半部分全选, 对应的状态的二进制表示是全1, 即111111...11, 这个状态是在最后一次迭代才会出现 
            if (tot != right_sum && left.count(-tot)) return true;
        }

        return false;
    }
};

思路三

动态规划 + 二进制优化

还是接前面的思路一，我们最终要找到一个数组A，满足sumA / k = sum / n。

转换一下，即我们需要选择k个数，使得它们的和等于 sum * k / n。

感觉有点像 0-1 背包，但是这里的条件是动态变化的。因为需要满足的条件，sumA = sum * k / n当中，k是动态变化的。而0-1背包的背包体积是固定不变的。但是我们依然可以用动态规划的思路来做。由于我们关注两个信息：选取的数的个数，以及这些数的和。

那么可以考虑用一个二维数组来存储状态，用dp[s][i]来表示，从数组中选取i个数，且这i个数的和为s，是否存在这样的方案。

若dp[s][i] = true表示，能够从数组中选出i个数，并且这些数的和为s。

那么状态如何转移呢？如果能选择i个数，使得和为s，即dp[s][i] = true，那么此时再从未选择的数中选一个数num，使得dp[s + num][i + 1] = true

所以转移方程是

if (dp[s][i] == true) {
    dp[s + num][i + 1] = true;
}

在计算出全部的状态数组后，我们只需要枚举选择的个数（k ∈[1, n - 1]），计算一下此时的sumA = sum * k / n，看一下dp[sumA][k] 是否为true即可，我们来写一下代码。

// C++
class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        int n = nums.size();
        if (n <= 1) return false;
        int sum = accumulate(nums.begin(), nums.end(), 0);
        // dp数组, 最大的数据范围是 dp[sum][n]
        vector<vector<bool>> dp(sum + 1, vector<bool>(n + 1));
        // 初始化边界状态
        dp[0][0] = true; // 一个数都不选时
        // 这里要特别注意状态转移的计算顺序!!!!
        // 最外层循环, 枚举nums, 每次看一下当前的数能否加到某个方案中
        // 如果最外层不枚举nums, 则会产生一个数被重复添加几次的情况
        for (int& num : nums) {
            // 这里的和, 需要从大到小枚举, 而不能从小到大枚举
            // 如果和从小到大枚举, 可能先更新了dp[num][j] = true, 然后随后又更新了 dp[num * 2][j] = true, 这样就把 num 添加了两次
            for (int s = sum; s >= num; s--) {
                for (int j = 1; j <= n; j++) {
                    if (dp[s - num][j - 1]) {
                        dp[s][j] = true;
                    }
                }
            }
        }

        for (int k = 1; k < n; k++) {
            if (sum * k % n) continue; // sumA = sum * k / n 不是整数, 则直接跳过
            if (dp[sum * k / n][k]) return true;
        }
        return false;
    }
};

这份代码提交后会报TLE，因为时间复杂度太高了。我们来看一下这样做的时间复杂度是多少，一共三层循环，时间复杂度为 $O(n^2 × sum)$ ，n = 30，sum = 30 × 10^4，大概在 $2.7×10^8$ ，一定会超时。

我们发现，n = 30，比较小，我们可以用二进制优化。具体怎么做呢？对于上面定义的二维dp数组的第一维，当第一维固定为num时，其第二维最多有30个值左右，即dp[num][0]，dp[num][1]，dp[num][2]，…，dp[num][30]，这每个状态，要么是true要么是false。那么我们可以将每个值用一个二进制位来表示，二进制位为1表示true，二进制位为0表示false。

比如对于num = 3，其状态为dp[3][0] = false，dp[3][1] = true，dp[3][2] = true，dp[3][3] = false

那么我们可以对num = 3的这4个状态，用一个二进制数来表示，即dp[3] = 0110。（第二维0，对应二进制表示的第0位，第二维1，对应二进制表示的第1位，…）

这样就将二维状态表示，通过二进制优化，压缩成了一维。这样，我们可以将内层循环中的[1, n]次运算，缩减为1次。

具体怎么做呢？由于我们只需要处理原先二维状态中dp[num][i] = true的那些i，那么只要一维状态的dp[num]不为0，则说明有一些二进制位为1，对于这些二进制位，我们可以将其往左移一位，则表示了个数+1。

注意，我们用的是二进制位所在的位置，来表示选择了多少个数，最低位从0开始，表示一个也不选

比如，一维状态``dp[num] = 0011，表示二维状态dp[num][0] = true，dp[num][1] = true`。

// C++ 188ms
class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        int n = nums.size();
        if (n <= 1) return false;
        int sum = accumulate(nums.begin(), nums.end(), 0);
        // dp数组, 最大的数据范围是 dp[sum][n]
        vector<int> dp(sum + 1);
        // 初始化边界状态
        dp[0] = 1; // 一个数都不选时, 即选择的个数为0, 即第0位的二进制位为1
        // 这里要特别注意状态转移的计算顺序!!!!
        // 最外层循环, 枚举nums, 每次看一下当前的数能否加到某个方案中
        // 如果最外层不枚举nums, 则会产生一个数被重复添加几次的情况
        for (int& num : nums) {
            // 这里的和, 需要从大到小枚举, 而不能从小到大枚举
            // 如果和从小到大枚举, 可能先更新了dp[num][j] = true, 然后随后又更新了 dp[num * 2][j] = true, 这样就把 num 添加了两次
            for (int s = sum; s >= num; s--) {
                if (dp[s - num]) {
                    // 这个状态不为0, 则说明有一些二进制位是1
                    // 将那些二进制位表示的个数, 全部+1, 再和dp[s] 做一下或运算
                    // 相当于一次计算了所有dp[s][0], dp[s][1], dp[s][2], dp[s][3], ......
                    dp[s] |= dp[s - num] << 1;
                }
            }
        }

        for (int k = 1; k < n; k++) {
            if (sum * k % n) continue; // sumA = sum * k / n 不是整数, 则直接跳过
            if ((dp[sum * k / n] >> k) & 1) return true; // 第k位为1时, 则有效
        }
        return false;
    }
};

另：上面的状态数组的第二维，其实可以不用枚举到n，只用枚举到n / 2。为什么呢？因为A组和B组互为镜像，我们只用找到一个即可。若存在一个有效的划分，其中A组的个数大于n / 2，那么B组的个数一定是小于n / 2的。所以我们只需要找到个数小于等于n / 2的一种方案即可，这算是一个小小的优化。

朴素动规：

// C++
class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        int n = nums.size();
        if (n <= 1) return false;
        int sum = accumulate(nums.begin(), nums.end(), 0);
        // dp数组, 最大的数据范围是 dp[sum][n]
        vector<vector<bool>> dp(sum + 1, vector<bool>(n / 2 + 1));
        // 初始化边界状态
        dp[0][0] = true; // 一个数都不选时
        // 这里要特别注意状态转移的计算顺序!!!!
        // 最外层循环, 枚举nums, 每次看一下当前的数能否加到某个方案中
        // 如果最外层不枚举nums, 则会产生一个数被重复添加几次的情况
        for (int& num : nums) {
            // 这里的和, 需要从大到小枚举, 而不能从小到大枚举
            // 如果和从小到大枚举, 可能先更新了dp[num][j] = true, 然后随后又更新了 dp[num * 2][j] = true, 这样就把 num 添加了两次
            for (int s = sum; s >= num; s--) {
                for (int j = 1; j <= n / 2; j++) {
                    if (dp[s - num][j - 1]) {
                        dp[s][j] = true;
                    }
                }
            }
        }

        for (int k = 1; k <= n / 2; k++) {
            if (sum * k % n) continue; // sumA = sum * k / n 不是整数, 则直接跳过
            if (dp[sum * k / n][k]) return true;
        }
        return false;
    }
};

动规+二进制优化

// C++ 188ms
class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        int n = nums.size();
        if (n <= 1) return false;
        int sum = accumulate(nums.begin(), nums.end(), 0);
        // dp数组, 最大的数据范围是 dp[sum][n]
        vector<int> dp(sum + 1);
        // 初始化边界状态
        dp[0] = 1; // 一个数都不选时, 即选择的个数为0, 即第0位的二进制位为1
        // 这里要特别注意状态转移的计算顺序!!!!
        // 最外层循环, 枚举nums, 每次看一下当前的数能否加到某个方案中
        // 如果最外层不枚举nums, 则会产生一个数被重复添加几次的情况
        for (int& num : nums) {
            // 这里的和, 需要从大到小枚举, 而不能从小到大枚举
            // 如果和从小到大枚举, 可能先更新了dp[num][j] = true, 然后随后又更新了 dp[num * 2][j] = true, 这样就把 num 添加了两次
            for (int s = sum; s >= num; s--) {
                if (dp[s - num]) {
                    // 这个状态不为0, 则说明有一些二进制位是1
                    // 将那些二进制位表示的个数, 全部+1, 再和dp[s] 做一下或运算
                    // 相当于一次计算了所有dp[s][0], dp[s][1], dp[s][2], dp[s][3], ......
                    dp[s] |= dp[s - num] << 1;
                }
            }
        }

        for (int k = 1; k <= n / 2; k++) {
            if (sum * k % n) continue; // sumA = sum * k / n 不是整数, 则直接跳过
            if ((dp[sum * k / n] >> k) & 1) return true; // 第k位为1时, 则有效
        }
        return false;
    }
};