LeetCode Solutions Blog

Partition to K Equal Sum Subsets

Number: 698

Difficulty: Medium

Paid? No

Problem Description

Given an integer array nums and an integer k, determine if it is possible to partition the array into k non-empty subsets whose sums are all equal. For example, given nums = [4,3,2,3,5,2,1] and k = 4, one valid partition is (5), (1,4), (2,3), (2,3), each with a sum of 5.

Key Insights

The total sum of the array must be divisible by k; otherwise, it’s impossible to partition it into subsets of equal sum.
The target sum for each subset is the total sum divided by k.
Sorting the array in descending order helps in pruning branches early in the backtracking process.
A backtracking (DFS) approach, optionally enhanced with memoization or bitmasking, can be used to explore possible subset partitions.
Handling duplicate numbers carefully reduces redundant work during recursion.

Space and Time Complexity

Time Complexity: O(k * 2^n) in the worst-case scenario (where n is the number of elements) due to exploring subset combinations via backtracking. Space Complexity: O(n) for recursion stack and additional data structures (or O(2^n) if using memoization with bitmasking).

Solution

We first check if the total sum of the array is divisible by k. Then determine the target sum for each subset. By sorting the array in descending order, we ensure that the larger numbers are placed first so that invalid configurations are pruned quickly. We use a backtracking (DFS) approach to try and assign numbers to each subset while ensuring the running sum does not exceed the target. When a subset reaches the target sum, we recursively try to fill the next subset. Care is taken to skip duplicates in order to avoid redundant recursion. This approach efficiently determines if a valid k-partition exists.

Code Solutions

# Define the class with the solution method
class Solution:
    def canPartitionKSubsets(self, nums, k):
        # Calculate the total sum and check divisibility by k
        total_sum = sum(nums)
        if total_sum % k != 0:
            return False
        
        target = total_sum // k
        
        # Sort the numbers in descending order for optimization
        nums.sort(reverse=True)
        # If the largest number is greater than the target, partition is impossible
        if nums[0] > target:
            return False
        
        n = len(nums)
        used = [False] * n
        
        # Backtracking function: start_index is where to continue searching, 
        # curr_sum is the current subset sum, count is how many subsets have been formed so far.
        def backtrack(start_index, curr_sum, count):
            # If we've correctly formed k-1 subsets, the rest of the numbers form the last subset.
            if count == k - 1:
                return True
            # When current subset sum equals target, start forming a new subset.
            if curr_sum == target:
                return backtrack(0, 0, count + 1)
            # Try to assign next numbers to current subset.
            for i in range(start_index, n):
                if not used[i] and curr_sum + nums[i] <= target:
                    used[i] = True
                    if backtrack(i + 1, curr_sum + nums[i], count):
                        return True
                    used[i] = False
                    # Skip duplicate elements to prevent redundant work.
                    while i + 1 < n and nums[i] == nums[i + 1]:
                        i += 1
            return False
        
        return backtrack(0, 0, 0)

# Example usage:
# sol = Solution()
# print(sol.canPartitionKSubsets([4,3,2,3,5,2,1], 4))  # Output: True