Medium

Partition to K Equal Sum Subsets

LAB

Return true if nums can be partitioned into k non-empty subsets with equal sum.

EXAMPLES

Example 1

Input
{
  "nums": [
    4,
    3,
    2,
    3,
    5,
    2,
    1
  ],
  "k": 4
}

Output
true

FUNCTION SHAPE

nums: intArrayk: int→bool

SOLUTION NOTE

Bitmask tracks used elements. curr_sum resets to 0 when reaching target (completed one subset).

Reveal reference solution +

pythonREFERENCE

def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
    total = sum(nums)
    if total % k != 0:
        return False
    target = total // k
    n = len(nums)

    @cache
    def dp(mask, curr_sum):
        if mask == (1 << n) - 1:
            return True
        for i in range(n):
            if not (mask & (1 << i)) and curr_sum + nums[i] <= target:
                next_sum = (curr_sum + nums[i]) % target
                if dp(mask | (1 << i), next_sum):
                    return True
        return False

    return dp(0, 0)

TimeO(n × 2^n)

SpaceO(2^n)

Open on LeetCode ↗