Partition to K Equal Sum Subsets

update 2018-10-10 00:53:15

LeetCode

Given an array of integers nums and a positive integer k, find whether it's 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's possible to divide it into 4 subsets (5), (1, 4), (2,3), (2,3) with equal sums.

Note:

1. 1 <= k <= len(nums) <= 16.

2. 0 < nums[i] < 10000.

Basic Idea:

这是 416. Partition Equal Subset Sum 的follow up，方法也和之前的题目类似。

初拿到题目的时候感觉很蒙逼，但是仔细想了一下，发现有如下隐含信息：

1. 要把给定数组分成k个sum相等的subset，可以推出targetSum一定为整数，并且等于 sum/k；

2. 如果发现sum/k的值不是整数，可以直接返回false;

这样问题就被转换成了：求给定数组能否被分成k个subset并且他们的sum都为targetSum。那么如何求解呢？想到可以使用dfs。这里因为subset的数量已经确定了，我们可以先虚拟创建出所有的subset，然后考虑每个数字被放入不同subset的情况。这种DFS的思路和之前的题目都是不同的，需要格外注意。

起初我的想法是先创建第一个subset，然后第二个。。以此类推，但是这样不容易实现，逻辑太复杂。最后的做法是创建一个长度为 k 的数组表示每个subset的sum值，然后进行dfs，每层考虑一个index的数字会被放入哪个subset，共有len(nums)层。时间复杂度为O(k^len(nums))。

优化思路主要有三点：

• 如果targetSum不为整数直接返回false；

• 如果存在 num>targetSum 直接返回 false；

• 如果在某一层有几个subset的当前sum相等，则只需要考虑其中一个；

• 可以对nums排序，然后从大到小考虑每个数字，这样可以提前发现行不通的路；

Java Code

class Solution {
    public boolean canPartitionKSubsets(int[] nums, int k) {
        int sum = 0;
        for (int num : nums) sum += num;
        if (sum % k != 0) return false; // 如果sum/k除不尽，说明不存在这样的partition
        sum /= k;
        for (int num : nums) if (num > sum) return false;
        int[] setList = new int[k];
        return helper(nums, setList, 0, sum);
    }

    private boolean helper(int[] nums, int[] setList, int pos, int targetSum) {
        if (pos == nums.length) return true;
        Set<Integer> currSumSet = new HashSet<>();
        for (int i = 0; i < setList.length; ++i) {
            int currSum = setList[i];
            if (! currSumSet.add(currSum)) continue;
            if (nums[pos] + currSum <= targetSum) {
                setList[i] += nums[pos];
                if (helper(nums, setList, pos + 1, targetSum)) return true;
                setList[i] -= nums[pos];
            }
        }
        return false;
    }
}

Update: 09/30/2021

之前的做法时间复杂度分析：

1. 每层考虑一个数字，分别尝试放入k个set

2. 一共有N个数字

3. 总时间复杂度则为 O(k^N)

之前的做法时间复杂度是比较大的，我们可以用另一种思路来考虑，即分别考虑每个subset，尝试生成一个subset之后再尝试下一个，直到生成了全部K个subsets。在下面的实现中，我们每次从0 index开始考虑每个数字要不要加入当前subset，对每个subset来讲，时间复杂度为 O(2^N), 一共有K个subset，所以一共为 O(K * 2^N) , 这是优于之前的方法的。

Java Code:

class Solution {
    public boolean canPartitionKSubsets(int[] nums, int k) {
        if (nums.length < k) return false;
        int sum = 0;
        for (int num : nums) sum += num;
        if (sum % k != 0) return false;
        int target = sum / k;
        for (int num : nums) if (num > target) return false;
        return dfs(new int[k], nums, 0, target);
    }
    
    private boolean dfs(int[] sets, int[] nums, int pos, int target) {
        if (pos == nums.length) return true;
        int num = nums[pos];
        boolean ret = false;
        Set<Integer> set = new HashSet<>();
        for (int i = 0; i < sets.length; ++i) {
            if (!set.add(sets[i])) continue;
            if (sets[i] + num <= target) {
                sets[i] += num;
                ret |= dfs(sets, nums, pos + 1, target);
                sets[i] -= num;
            }
        }
        return ret;
    }
}