力扣 698. 划分为k个相等的子集

题目

给定一个整数数组 nums 和一个正整数 k，找出是否有可能把这个数组分成 k 个非空子集，其总和都相等。

示例

输入： nums = [4, 3, 2, 3, 5, 2, 1], k = 4
输出： True
说明： 有可能将其分成 4 个子集（5），（1,4），（2,3），（2,3）等于总和。

输入: nums = [1,2,3,4], k = 3
输出: false

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/partition-to-k-equal-sum-subsets
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

方法1：回溯

• 剪枝1：如果当前桶和上一个桶内的元素和相等，那么 nums[index] 选择上一个桶和选择当前桶可以得到的结果是一致的。
• 剪枝2：提前对 nums 数组排序，把大的数字排在前面，那么大的数字会先被分配到 bucket 中，对于之后的数字，bucket[i] + nums[index] 会更大，更容易触发剪枝的 if 条件。

Java实现

class Solution {
    public boolean canPartitionKSubsets(int[] nums, int k) {
        int sum = 0;
        for (int num : nums) sum += num;
        if (sum % k != 0) return false;

        Arrays.sort(nums);
        int l = 0, r = nums.length - 1;
        while (l < r) {
            int tmp = nums[l];
            nums[l] = nums[r];
            nums[r] = tmp;
            l++;
            r--;
        }

        int[] bucket = new int[k];
        int target = sum / k;
        return track_back(nums, 0, bucket, k, target);
    }

    public boolean track_back(int[] nums, int idx, int[] bucket, int k, int target){
        if (idx == nums.length) return true;

        for (int i = 0; i < k; i++) {
            if (i > 0 && bucket[i - 1] == bucket[i]) continue; //剪枝1
            
            if (nums[idx] + bucket[i] > target) continue; //剪枝2
            
            bucket[i] += nums[idx];
            if (track_back(nums, idx + 1, bucket, k, target)) return true;
            bucket[i] -= nums[idx];
        }
        return false;
    }
}