python https://leetcode.com/problems/partition-to-k-equal-sum-subsets/description/

Posted 2021-05-09

# Time: O(k*2^n)
# Space: O(k + n) where n = len(nums)
# 698. Partition to K Equal Sum Subsets

class Solution(object):
    def canPartitionKSubsets(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: bool
        """
        # 1. Edge Case 
        if len(nums) < k or sum(nums) % k != 0:
            return False 
        subsum = sum(nums) / k
        for i in nums:
            if i > subsum:
                return False 
        # 2. Preprocess
        parts = [0] * k
        taken = [False] * len(nums)
        return self.backtrack(0, 0, nums, k, subsum, parts, taken)

    # 3. Backtrack
    def backtrack(self, p, t, nums, k, subsum, parts, taken):
        # 3.1 Base Case 
        if p == k:
            return True 
        elif parts[p] == subsum:
            return self.backtrack(p+1, 0, nums, k, subsum, parts, taken)
        # 3.2 to take or not 
        for i in range(t, len(nums)): # !!! start from t, otherwise TLE 
            if not taken[i] and parts[p] + nums[i] <= subsum:
                taken[i] = True 
                parts[p] += nums[i]
                if self.backtrack(p, i+1, nums, k, subsum, parts, taken):
                    return True
                parts[p] -= nums[i]
                taken[i] = False 
        return False