给定一个整数数组
nums
和一个正整数k
,找出是否有可能把这个数组分成k
个非空子集,其总和都相等。示例 1:
输入: nums = [4, 3, 2, 3, 5, 2, 1], k = 4 输出: True 说明: 有可能将其分成 4 个子集(5),(1,4),(2,3),(2,3)等于总和。示例 2:
输入: nums = [1,2,3,4], k = 3 输出: false
刚开始读这道题的时候没有一个好的思路那就暴力呗。
如何暴力?有两种方式:把k个非空子集看成是k个盒子最终要保证把装有 n 个数字的数组 nums 分成 k 个和相同的集合,这些盒子中存储数值的总和保持一致。然后分成两种情况:
1. 对于nums中的每个数选择进入到k个盒子中的某一个
2.对于每个盒子,遍历nums中的n个数字,选择是否将当前遍历的数字放入盒子中
他们的主要区别在于时间、空间复杂度的不同,我们要分析比较选出复杂度更低的解法。
好了,让我们先来看一下第一种情况:
// k 个盒子(集合),记录每个盒子装的数字之和
int[] bucket = new int[k];
// 穷举 nums 中的每个数字
for (int index = 0; index < nums.length; index++) {
// 穷举每个盒子
for (int i = 0; i < k; i++) {
// nums[index] 选择是否要进入第 i 个盒子
// 选择装进第 i 个盒子
bucket[i] += nums[index];
// 递归穷举下一个数字的选择
backtrack(nums, index + 1);
// 撤销选择
bucket[i] -= nums[index];
}
}
这样子提交是TLE的那就得再想想如何优化了——剪枝
如何剪枝?
if (bucket[i] + nums[index] > target) {
continue;
}
这样子还是TLE还需要进一步优化
如果我们让尽可能多的情况命中剪枝的那个 if 分支,就可以减少递归调用的次数,一定程度上减少时间复杂度。
如何尽可能多的命中这个 if 分支呢?要知道我们的 index 参数是从 0 开始递增的,也就是递归地从 0 开始遍历 nums 数组。
如果我们提前对 nums 数组排序,把大的数字排在前面,那么大的数字会先被分配到 bucket 中,对于之后的数字,bucket[i] + nums[index] 会更大,更容易触发剪枝的 if 条件。
boolean canPartitionKSubsets(int[] nums, int k) {
// 其他代码不变
// ...
/* 降序排序 nums 数组 /
Arrays.sort(nums);
// 翻转数组,得到降序数组
for (i = 0, j = nums.length - 1; i < j; i++, j--) {
int temp = nums[i];
nums[i] = nums[j];
nums[j] = temp;
}
* /*******************/
return backtrack(nums, 0, bucket, target);
}
完整代码:
class Solution {
public boolean canPartitionKSubsets(int[] nums, int k) {
if (k > nums.length) return false;
int sum = 0;
for (int v : nums) sum += v;
if (sum % k != 0) return false;
int[] bucket = new int[k];
int target = sum / k;
Arrays.sort(nums);
for (int i = 0, j = nums.length - 1; i < j; i++ , j--) {
int tmp = nums[i];
nums[i] = nums[j];
nums[j] = tmp;
}
return backtrack(nums, 0, bucket, target);
}
boolean backtrack(int[] nums, int index, int[] bucket, int target) {
if (index == nums.length) {
for (int i = 0; i < bucket.length; i++ ) {
if (bucket[i] != target) return false;
}
return true;
}
for (int i = 0; i < bucket.length; i++ ) {
if (bucket[i] + nums[index] > target) continue;
bucket[i] += nums[index];
if (backtrack(nums, index + 1, bucket, target)) {
return true;
}
bucket[i] -= nums[index];
}
return false;
}
}
但是这样还是TLE,我真崩溃了TAT
别急啊这只是讨论了一种视角,还有另一种视角呢
盒子视角:
以盒子的视角进行穷举,每个盒子需要遍历 nums 中的所有数字,决定是否把当前数字装进盒子中;当装满一个盒子之后,还要装下一个盒子,直到所有盒子都装满为止。
// 装满所有盒子为止
while (k > 0) {
// 记录当前盒子中的数字之和
int bucket = 0;
for (int i = 0; i < nums.length; i++) {
// 决定是否将 nums[i] 放入当前盒子中
if (canAdd(bucket, num[i])) {
bucket += nums[i];
}
if (bucket == target) {
// 装满了一个盒子,装下一个盒子
k--;
break;
}
}
}
也可以把这个 while 循环改写成递归函数,不过比刚才略微复杂一些,首先写一个 backtrack 递归函数出来:
boolean backtrack(int k, int bucket, int[] nums, int start, boolean[] used, int target);
他的含义是:
现在 k 号盒子正在思考是否应该把 nums[start] 这个元素装进来;目前 k 号盒子里面已经装的数字之和为 bucket;used 标志某一个元素是否已经被装到盒子中;target 是每个盒子需要达成的目标和。
步骤:
1.需要遍历 nums 中所有数字,决定哪些数字需要装到当前桶中。
2.如果当前桶装满了(桶内数字和达到 target),则让下一个桶开始执行第 1 步。
boolean backtrack(int k, int bucket, int[] nums, int start, boolean[] used, int target) {
// base case
if (k == 0) {
// 所有盒子都被装满了,而且 nums 一定全部用完了
// 因为 target == sum / k
return true;
}
if (bucket == target) {
// 装满了当前盒子,递归穷举下一个盒子的选择
// 让下一个盒子从 nums[0] 开始选数字
return backtrack(k - 1, 0 ,nums, 0, used, target);
}
// 从 start 开始向后探查有效的 nums[i] 装入当前盒子
for (int i = start; i < nums.length; i++) {
// 剪枝
if (used[i]) {
// nums[i] 已经被装入别的盒子中
continue;
}
if (nums[i] + bucket > target) {
// 当前盒子装不下 nums[i]
continue;
}
// 做选择,将 nums[i] 装入当前盒子中
used[i] = true;
bucket += nums[i];
// 递归穷举下一个数字是否装入当前盒子
if (backtrack(k, bucket, nums, i + 1, used, target)) {
return true;
}
// 撤销选择
used[i] = false;
bucket -= nums[i];
}
// 穷举了所有数字,都无法装满当前盒子
return false;
}
我能想到的也就到这里了,下面是东哥的原文:
这段代码是可以得出正确答案的,但是效率很低,我们可以思考一下是否还有优化的空间。
首先,在这个解法中每个桶都可以认为是没有差异的,但是我们的回溯算法却会对它们区别对待,这里就会出现重复计算的情况。
什么意思呢?我们的回溯算法,说到底就是穷举所有可能的组合,然后看是否能找出和为 target 的 k 个桶(子集)。
那么,比如下面这种情况,target = 5,算法会在第一个桶里面装 1, 4:
现在第一个桶装满了,就开始装第二个桶,算法会装入 2, 3:
然后以此类推,对后面的元素进行穷举,凑出若干个和为 5 的桶(子集)。
但问题是,如果最后发现无法凑出和为 target 的 k 个子集,算法会怎么做?
回溯算法会回溯到第一个桶,重新开始穷举,现在它知道第一个桶里装 1, 4 是不可行的,它会尝试把 2, 3 装到第一个桶里:
现在第一个桶装满了,就开始装第二个桶,算法会装入 1, 4:
好,到这里你应该看出来问题了,这种情况其实和之前的那种情况是一样的。也就是说,到这里你其实已经知道不需要再穷举了,必然凑不出来和为 target 的 k 个子集。
但我们的算法还是会傻乎乎地继续穷举,因为在她看来,第一个桶和第二个桶里面装的元素不一样,那这就是两种不一样的情况呀。
那么我们怎么让算法的智商提高,识别出这种情况,避免冗余计算呢?
你注意这两种情况的 used 数组肯定长得一样,所以 used 数组可以认为是回溯过程中的「状态」。
所以,我们可以用一个 memo 备忘录,在装满一个桶时记录当前 used 的状态,如果当前 used 的状态是曾经出现过的,那就不用再继续穷举,从而起到剪枝避免冗余计算的作用。
有读者肯定会问,used 是一个布尔数组,怎么作为键进行存储呢?这其实是小问题,比如我们可以把数组转化成字符串,这样就可以作为哈希表的键进行存储了。
class Solution {
// 备忘录,存储 used 数组的状态
HashMap<String, Boolean> memo = new HashMap<>();
public boolean canPartitionKSubsets(int[] nums, int k) {
// 见上文
}
boolean backtrack(int k, int bucket, int[] nums, int start, boolean[] used, int target) {
// base case
if (k == 0) {
return true;
}
// 将 used 的状态转化成形如 [true, false, ...] 的字符串
// 便于存入 HashMap
String state = Arrays.toString(used);
if (bucket == target) {
// 装满了当前桶,递归穷举下一个桶的选择
boolean res = backtrack(k - 1, 0, nums, 0, used, target);
// 将当前状态和结果存入备忘录
memo.put(state, res);
return res;
}
if (memo.containsKey(state)) {
// 如果当前状态曾今计算过,就直接返回,不要再递归穷举了
return memo.get(state);
}
// 其他逻辑不变...
}
}
这不就到了记忆搜索的知识点了嘛!
可以去看我之前的博客有讲过记忆搜索
这样提交解法,发现执行效率依然比较低,这次不是因为算法逻辑上的冗余计算,而是代码实现上的问题。
因为每次递归都要把 used 数组转化成字符串,这对于编程语言来说也是一个不小的消耗,所以我们还可以进一步优化。
注意题目给的数据规模 nums.length <= 16,也就是说 used 数组最多也不会超过 16,那么我们完全可以用「位图」的技巧,用一个 int 类型的 used 变量来替代 used 数组。
具体来说,我们可以用整数 used 的第 i 位((used >> i) & 1)的 1/0 来表示 used[i] 的 true/false。
这样一来,不仅节约了空间,而且整数 used 也可以直接作为键存入 HashMap,省去数组转字符串的消耗。
class Solution {
public boolean canPartitionKSubsets(int[] nums, int k) {
// 排除一些基本情况
if (k > nums.length) return false;
int sum = 0;
for (int v : nums) sum += v;
if (sum % k != 0) return false;
int used = 0; // 使用位图技巧
int target = sum / k;
// k 号桶初始什么都没装,从 nums[0] 开始做选择
return backtrack(k, 0, nums, 0, used, target);
}
HashMap<Integer, Boolean> memo = new HashMap<>();
boolean backtrack(int k, int bucket,
int[] nums, int start, int used, int target) {
// base case
if (k == 0) {
// 所有桶都被装满了,而且 nums 一定全部用完了
return true;
}
if (bucket == target) {
// 装满了当前桶,递归穷举下一个桶的选择
// 让下一个桶从 nums[0] 开始选数字
boolean res = backtrack(k - 1, 0, nums, 0, used, target);
// 缓存结果
memo.put(used, res);
return res;
}
if (memo.containsKey(used)) {
// 避免冗余计算
return memo.get(used);
}
for (int i = start; i < nums.length; i++) {
// 剪枝
if (((used >> i) & 1) == 1) { // 判断第 i 位是否是 1
// nums[i] 已经被装入别的桶中
continue;
}
if (nums[i] + bucket > target) {
continue;
}
// 做选择
used |= 1 << i; // 将第 i 位置为 1
bucket += nums[i];
// 递归穷举下一个数字是否装入当前桶
if (backtrack(k, bucket, nums, i + 1, used, target)) {
return true;
}
// 撤销选择
used ^= 1 << i; // 使用异或运算将第 i 位恢复 0
}
}
为什么第一种解法即便经过了排序优化,也明显比第二种解法慢很多,这是为什么呢?
我们来分析一下这两个算法的时间复杂度,假设 nums 中的元素个数为 n。
先说第一个解法,也就是从数字的角度进行穷举,n 个数字,每个数字有 k 个桶可供选择,所以组合出的结果个数为 k^n,时间复杂度也就是 O(k^n)。
第二个解法,每个桶要遍历 n 个数字,对每个数字有「装入」或「不装入」两种选择,所以组合的结果有 2^n 种;而我们有 k 个桶,所以总的时间复杂度为 O(k*2^n)。
当然,这是对最坏复杂度上界的粗略估算,实际的复杂度肯定要好很多,毕竟我们添加了这么多剪枝逻辑。不过,从复杂度的上界已经可以看出第一种思路要慢很多了。
所以,谁说回溯算法没有技巧性的?虽然回溯算法就是暴力穷举,但穷举也分聪明的穷举方式和低效的穷举方式,关键看你以谁的「视角」进行穷举。
通俗来说,我们应该尽量「少量多次」,就是说宁可多做几次选择(乘法关系),也不要给太大的选择空间(指数关系);做 n 次「k 选一」仅重复一次(O(k^n)),比 n 次「二选一」重复 k 次(O(k*2^n))效率低很多。
师承我东哥