Leetcode 698 划分为k个相等的子集 ( Partition to K Equal Sum Subsets *Medium* ) 题解分析

题目介绍

Given an integer array nums and an integer k, return true if it is possible to divide this array into k non-empty subsets whose sums are all equal.

示例

Example 1:

Input: nums = [4,3,2,3,5,2,1], k = 4
Output: true
Explanation: It is possible to divide it into 4 subsets (5), (1, 4), (2,3), (2,3) with equal sums.

Example 2:

Input: nums = [1,2,3,4], k = 3
Output: false

Constraints:

  • 1 <= k <= nums.length <= 16
  • 1 <= nums[i] <= 10^4
  • The frequency of each element is in the range [1, 4].

解析

看到这个题一开始以为挺简单,但是仔细想想问题还是挺多的,首先是分成 k 组,但是数量不限,应该需要用到回溯的方式,同时对于时间和空间复杂度也有要求,一开始这个代码是超时的,我也试了下 leetcode 上 discussion 里 vote 最高的提交也是超时的,不过看 discussion 里的帖子,貌似是后面加了一些条件,可以帮忙提高执行效率,第三条提示不太清楚意图,具体可以看下代码

代码

public boolean canPartitionKSubsets(int[] nums, int k) {
    if (k == 1) {
        return true;
    }
    int sum = 0, n;
    n = nums.length;
    for (int num : nums) {
        sum += num;
    }
    if (sum % k != 0) {
        return false;
    }

    int avg = sum / k;
    // 排序
    Arrays.sort(nums);
    // 做个前置判断,如果最大值超过分组平均值了就可以返回 false 了
    if (nums[n - 1] > avg) {
        return false;
    }
    // 这里取了个巧,先将数组中元素就等于分组平均值的直接排除了
    int calculated = 0;
    for (int i = n - 1; i > 0; i--) {
        if (nums[i] == avg) {
            k--;
            calculated++;
        }
    }

    int[] bucket = new int[k];
    // 初始化 bucket
    for (int i = 0; i < k; i++) {
        bucket[i] = avg;
    }

    // 提前做下边界判断
    if (nums[n - 1] > avg) {
        return false;
    }

    return backTraversal(nums, bucket, k, n - 1 - calculated);
}

private boolean backTraversal(int[] nums, int[] bucket, int k, int cur) {
    if (cur < 0) {
        return true;
    }
    for (int i = 0; i < k; i++) {
        if (bucket[i] == nums[cur] || bucket[i] >= nums[cur] + nums[0]) {
            // 判断如果当前 bucket[i] 剩余的数字等于nums[cur], 即当前bucket已经满了
            // 或者如果当前 bucket[i] 剩余的数字大于等于 nums[cur] + nums[0] ,
            // 因为nums 在经过排序后 nums[0]是最小值,如果加上 nums[0] 都已经超过bucket[i] 了,
            // 那当前bucket[i] 肯定是没法由包含 nums[cur] 的组合组成一个满足和为前面 s/k 的组合了
            // 这里判断的是 nums[cur] ,如果第一次 k 次循环都不符合其实就返回 false 了

            // 而如果符合,就将 bucket[i] 减去 nums[cur] 再次进入递归,
            // 这里进入递归有个收敛参数就是 cur - 1,因为其实判断 cur 递减作为一个结束条件
            bucket[i] -= nums[cur];
            // 符合条件,这里对应着入口,当 cur 被减到 0 了,就表示都符合了因为是根据所有值的和 s 和 k 组除出来的平均值,当所有数都通过前面的 if 判断符合了,并且每个数字都使用了,
            // 即说明已经符合要求了
            if (backTraversal(nums, bucket, k, cur - 1)) return true;
            // 这边是个回退机制,如果前面 nums[cur]没办法组合成和为平均值的话就减掉进入下一个循环
            bucket[i] += nums[cur];
        }
    }
    return false;
}

最后贴个图