動機

人生第一次用狀態壓縮

Problem

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 = 4Output: trueExplanation: It's 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 = 3Output: false

 

Constraints:

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

Sol

去湊,湊完就繼續往下再湊

class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        cnt = sum(nums)
        if cnt % k != 0:
            return False
        else:
            cnt = cnt // k
            @cache
            def dp(tbl,acc=0):
                if tbl == 0:
                    return (acc == cnt)
                elif acc > cnt:
                    return False
                elif acc == cnt:
                    return dp(tbl,0)
                else:
                    cs = [i for i in range(len(nums)) if (tbl & (1 << i)) != 0]
                    return any([dp(tbl & ~(1 << i), acc+nums[i]) for i in cs])
            return dp(2**(len(nums))-1)