划分为k个相等的子集

标签: 位运算 记忆化搜索 数组 动态规划 回溯 状态压缩

难度: Medium

给定一个整数数组  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] 范围内

Submission

运行时间: 51 ms

内存: 16.0 MB

class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        def dfs(i):
            if i == len(nums):
                return True
            for j in range(k):
                if j and cur[j] == cur[j - 1]:
                    continue
                cur[j] += nums[i]
                if cur[j] <= s and dfs(i + 1):
                    return True
                cur[j] -= nums[i]
            return False

        s, mod = divmod(sum(nums), k)
        if mod:
            return False
        cur = [0] * k
        nums.sort(reverse=True)
        return dfs(0)

Explain

这个题解使用了回溯算法和剪枝优化。首先计算数组元素总和是否能被k整除,如果不能则直接返回False。然后将数组按从大到小排序,便于剪枝。接着使用递归回溯,维护一个大小为k的数组cur,表示当前k个子集的和。对于每个数,依次尝试放入每个子集中,如果放入后子集和不超过目标和s,则递归到下一个数。如果递归到最后一个数,说明找到了一种合法方案,返回True。如果尝试了所有可能性都无法找到合法方案,则返回False。在递归过程中,如果当前子集和与前一个子集和相等,则可以跳过,因为这种情况已经被之前的分支考虑过了,这样可以去除重复计算,减少时间复杂度。

时间复杂度: O(k^n)

空间复杂度: O(n)

```python
class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        def dfs(i):
            if i == len(nums):  # 已经遍历完所有数
                return True
            for j in range(k):  # 遍历k个子集
                if j and cur[j] == cur[j - 1]:  # 当前子集和与前一个相等,跳过
                    continue
                cur[j] += nums[i]  # 将当前数放入当前子集
                if cur[j] <= s and dfs(i + 1):  # 如果子集和不超过目标,递归到下一个数
                    return True
                cur[j] -= nums[i]  # 回溯,移除当前数
            return False

        s, mod = divmod(sum(nums), k)  
        if mod:  # 总和不能被k整除,返回False
            return False
        cur = [0] * k  # 初始化k个子集和为0
        nums.sort(reverse=True)  # 降序排列,便于剪枝
        return dfs(0)  # 从第0个数开始回溯
```

Explore

是的,考虑了所有元素都为零的情况。在这种特殊情况下,数组总和为零,如果k不为零,则总和除以k的余数也为零,因此可以被k整除。由于所有元素都是0,可以很容易地将它们分配到k个子集中,每个子集的和也将为0,满足题目要求。

在递归函数中,检查`if j and cur[j] == cur[j - 1]`是为了避免重复的子集配置产生相同的搜索路径。当两个连续的子集目前的和相同时,在添加新元素后,这两个子集是可交换的,即它们会产生相同的子集分配但顺序不同,从而导致重复计算。通过跳过这种情况,我们可以有效减少搜索空间,避免不必要的计算,提高算法效率。

数组进行降序排序后,可以帮助更快地达到子集和的上限。具体来说,因为较大的数先被处理,它们更快地填满接近目标和s的子集。如果一个较大的数无法放入任何当前的子集中(因为加入后会超过目标和s),那么较小的数也不可能放入,这样就可以立即中断当前的搜索路径。这种方法减少了不必要的递归调用,通过快速检测不可能的路径来提高效率。

是的,如果当前数加入任何子集后使得子集总和超过目标和s,则当前的分支会直接失败。在该情况下,算法会停止进一步探索这个分支并进行回溯,尝试将当前数加入其他子集或更改之前的分配。这是回溯法的一个关键特性,允许它通过尝试不同的组合来排除不可能满足条件的路径,从而找到可能的解决方案。