Leetcode 698. 划分为k个相等的子集 / 416. 分割等和子集 / 划分总和相同的子集 / python实现

题目

(473. 火柴拼正方形也是这一类型的，划分为四个子集的情况)

解题思路

题416是题698（k=2时）的特例。

划分为k个相等的子集的思路：
直接DFS判断是否能够划分成为k个子集，重点在于剪枝，即需要将数组从大到小进行排序。为什么不是从小到大排序？也是可以的，但是（在测试用例中）从大到小可以减少回溯的次数，大值+大值+… >= 目标值的几率更大。

当只划分为了两个集合时，可以直接套用上面划分为k个相等子集的代码，也可以利用01背包动态规划。动态规划思路：只要当前数组中能填充背包容量为sum(nums) // 2，那就可以划分为两个。填充背包相当于我们已经找到了一个子集，剩下的没有填充到背包中的元素总和肯定 = sum(nums) // 2，相当于另一个子集。

---

代码实现

"""
698. 划分为k个相等的子集
"""
class Solution:
    def canPartitionKSubsets(self, nums, k: int) -> bool:
        # 从大到小排序
        nums = sorted(nums, reverse=True)
        # 每个集合的目标总和
        sums = sum(nums) // k
        if sums < max(nums) or sum(nums) % k:
            return False
        vis = set()
        # 其中total表示当前集合的总和，times表示已经完成的集合数量
        def DFS(total, times):
            if total == sums:
                times += 1
                total = 0
            if not total and times == k and len(vis) == len(nums):
                return True
            for i in range(len(nums)):
                # （也是一个剪枝的小tip）相同的元素，之前没用上现在肯定也用不上
                if i and nums[i] == nums[i-1] and i-1 not in vis:
                    continue
                if i not in vis and total + nums[i] <= sums:
                    vis.add(i)
                    if DFS(total + nums[i], times):
                        return True
                    vis.remove(i)
            return False
        return DFS(0, 0)

"""
416. 分割等和子集 - 动态规划
"""
class Solution:
    def canPartition(self, nums) -> bool:
        sums = sum(nums) // 2
        if sums < max(nums) or sum(nums) % 2:
            return False
        # dp[i]表示数组是否能够填充背包容量为i
        dp = [False] * (sums+1)
        # 当背包容量为0的时候不需要填充，肯定可以实现
        dp[0] = True
        for i in range(len(nums)):
            for j in range(sums, -1, -1):
                if j < nums[i]:
                    continue
                # 该元素填入or不填入背包
                dp[j] |= dp[j - nums[i]]
        return dp[-1]

"""
416. 分割等和子集 - 回溯解法
"""
from functools import lru_cache
class Solution:
    def canPartition(self, nums) -> bool:
        sums = sum(nums) // 2
        if sums < max(nums) or sum(nums) % 2:
            return False
        vis = set()
        # 从大到小排序，减少回溯次数
        nums = sorted(nums, reverse=True)
        @lru_cache(None)
        # total表示当前子集的总和，other表示另一个子集的总和
        def DFS(total, other):
            if total == sums:
                return other == sums
            for i in range(len(nums)):
                # 剪枝，前一个没被选上，相同的元素也不会被选上
                if i and nums[i] == nums[i-1] and i-1 not in vis:
                    continue
                if i not in vis and total + nums[i] <= sums:
                    vis.add(i)
                    if DFS(total+nums[i],other-nums[i]):
                        return True
                    vis.remove(i)
            return False
        return DFS(0, sum(nums))