Partition to K Equal Sum Subsets

January 27, 2024

Problem Statement #

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’s possible to divide it into 4 subsets (5), (1, 4), (2, 3), (2, 3) with equal sums.

Solution Approach #

The solution involves using depth-first search (DFS) and backtracking to try and create k subsets with equal sums.

Algorithm Steps #

1. Calculate the sum of the array and check if it’s divisible by k. If not, return false.
2. Use an array to track which numbers have been used in a subset.
3. Recursively try to build each subset from the remaining numbers.
4. Use backtracking to undo choices that don’t lead to a solution.
5. Return true if k subsets with the target sum can be formed.

Code (Python) #

def canPartitionKSubsets(nums, k):
    total_sum = sum(nums)
    if total_sum % k != 0:
        return False
    target = total_sum // k
    used = [False] * len(nums)

    def backtrack(start, k, current_sum):
        if k == 0:
            return True
        if current_sum == target:
            return backtrack(0, k-1, 0)

        for i in range(start, len(nums)):
            if not used[i] and current_sum + nums[i] <= target:
                used[i] = True
                if backtrack(i+1, k, current_sum + nums[i]):
                    return True
                used[i] = False

        return False

    return backtrack(0, k, 0)

# Test the function
print(canPartitionKSubsets([4, 3, 2, 3, 5, 2, 1], 4))  # Output: True

Time Complexity #

• Worst case: O(k * 2^n), where n is the number of elements in nums. Each element can be in or out of a subset, and this is done for k subsets.

Space Complexity #

• O(n), for the recursion stack and the used array.