划分为K个相等的子集

2024-08-25

题目:

给定一个整数数组 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

提示:

  • 1 <= k <= len(nums) <= 16
  • 0 < nums[i] < 10000
  • 每个元素的频率在 [1,4] 范围内

思路:

如果这题是划分成两个和相等的子集,就非常好做了:
	先求原数组的总和 sum 
	然后通过回溯算法,在原数组里面找寻和为 sum / 2 的子集
	由于 2 = 1 + 1,找到 1 个,另一个自动产生

那对于本题的 k 个子集来说,就需要在原数组里面找 k 个和为 sum / k 的子集
	那么就需要回溯算法有一个参数,记录已经找到了几个满足要求的子集
	由于用过的元素不能再用,所以需要一个参数记录那些元素用过
		由于相同元素可以出现多次,所以不能用哈希表
		由于原数组固定,所以可以利用位数组,如果为1就表明该索引被使用
		位数字可以用一个数字表示,通过位运算操作使用和撤销
	剩下就是回溯算法的标准思路,从 0 开始寻找,判断是否可以加入子集中了

代码:

class Solution {
public:
    unordered_map<int, bool> memo; //备忘录

    bool backtrack(int k, int bucket, vector<int>& nums, int start, int used, int target) {
        if (k == 0) {
            // 所有桶都被装满了,而且 nums 一定全部用完了
            return true;
        }
        if (bucket == target) {
            // 装满了当前桶,递归穷举下一个桶的选择
            // 让下一个桶从 nums[0] 开始选数字
            bool res = backtrack(k - 1, 0, nums, 0, used, target);
            // 缓存结果
            memo[used] = res;
            return res;
        }
        if (memo.count(used)) {
            // 避免冗余计算
            return memo[used];
        }
        for (int i = start; i < nums.size(); 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
            bucket -= nums[i];
        }
        return false;
    }

    bool canPartitionKSubsets(vector<int>& nums, int k) {
        // 排除一些基本情况
        if (k > nums.size()) 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);
    }
};