Skip to content

698. Partition to K Equal Sum Subsets

Difficulty Topics

Description

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 = 4
Output: true
Explanation: It is 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 = 3
Output: false

 

Constraints:

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

Solution

partition-to-k-equal-sum-subsets.py
class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        n = len(nums)
        total = sum(nums)
        if total % k != 0: return False

        target = total // k
        nums.sort(reverse = 1)

        @cache
        def go(mask):
            curr = 0

            for i in range(n):
                if (mask >> i) & 1 > 0:
                    curr += nums[i]

            done, side = divmod(curr, target)

            if done == k - 1: return True

            for i in range(n):
                if (mask & (1 << i)) == 0:
                    if side + nums[i] <= target and go(mask | (1 << i)):
                        return True

            return False

        return go(0)