Partition to K Equal Sum Subsets

Patrick Leaton
Published: July 24, 2021, 8:35 p.m.
Problem 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's 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] <= 10^4
  • The frequency of each element is in the range [1, 4].

 

The description was taken from https://leetcode.com/problems/partition-to-k-equal-sum-subsets/.

Problem Solution

#O(K*2^N) Time, O(N) Space
class Solution(object):
    def canPartitionKSubsets(self, nums, k):
        if sum(nums)%k != 0:
            return False
 
        target = sum(nums)//k
 
        nums.sort()
        if nums[-1] > target:
            return False
 
        def backtrack(buckets):
            if not nums:
                return True
 
            candidate = nums.pop()
 
            for i, bucket in enumerate(buckets):
                if bucket + candidate <= target:
                    buckets[i] += candidate
                    if backtrack(buckets):
                        return True
                    buckets[i] -= candidate
 
                if bucket == 0:
                    break
       
            nums.append(candidate)
            return False
 
        return backtrack([0] * k)

Problem Explanation

To solve this question, we would need to find out if we can rearrange these matchsticks into k equal sum partitions where no number is split between multiple partitions.

This approach is a popular backtracking problem.

What we can do is initialize an array of k indices that we can picture as buckets.

Within each DFS path, we will pop a number off of the input array and place it into one of the buckets that have a sum less than or equal to the target value.  Then, we will continue down the DFS path keeping the current combination array.  If we have gotten to a point where we have placed each number, we know we have and can partition the matchsticks accordingly.  If not, we'll backtrack our placement choices and cycle the numbers back in to try again using different placement combinations.

First off, if we have leftover values after dividing the input array sum by the number of subsets we are trying to make, we can't have an equal sum, so we will return false.

        if sum(nums)%k != 0:
            return False

 

If that isn't the case, we will divide the input array sum by k.

This will give us the target value that each bucket needs to sum to.

        target = sum(nums)//k

 

Now, let's go ahead and sort the array so we can check if the last number within it is greater than the target value.  If it is, we couldn't possibly place that number into a subset and have it sum to the target value.  

If that's the case, we will return false.

        nums.sort()
        if nums[-1] > target:
            return False

Next, let's build our recursive, backtrack function.

This function will require a buckets variable.  This will basically be a combination array of size k that contains the sums for each subset in the current DFS path.

        def backtrack(buckets):

 

Let's draw out what we are hoping to achieve with this function.

If we are given the input where nums = [4,3,2,3,5,2,1] and k = 4, the total sum of nums would be twenty.

In this case, we would have four subsets that each sum to five.  We would end up with the buckets [5, 5, 5, 5], four values that signify four partitioned subsets of which the sums are equal.

We would know that we achieved this by whether or not we used up all of the numbers from our input array.  That is our base case.

            if not nums:
                return True

 

If we haven't, we will pop a number off the end of the input array and place that number into a bucket, if that number plus the value of the current bucket is less than equal to the target sum value.  This is so that we try each bucket that his number can fit into until our answer is found.

            candidate = nums.pop()
            for i, bucket in enumerate(buckets):
                if bucket + candidate <= target:
                    buckets[i] += candidate

 

Now we have chosen to put a number into one of the buckets, but we don't know if we made the right choice by placing it into whichever bucket it is now in.  To see, we will continue our DFS by using the current buckets combination and test whether we can place the remaining numbers into the remaining buckets with values less than the target.

We will know if this was the right choice if we hit our base case.

                    if backtrack(buckets):
                        return True
 

If not, we will backtrack the choice of placing the candidate number into the current bucket and try to place it in a different bucket later on.

                    buckets[i] -= candidate

 

To avoid endless loops, we will need to see if after removing the current candidate from the current bucket, the bucket is now zero.

If it is, that means we were unsuccessful in our placings of this bucket and don't want to keep trying to place the same combination of numbers.  If we put a number in and it didn't work, we don't want to put that same number in and get the same result, so we will break from this bucket.

                if bucket == 0:
                    break

 

If we have iterated through each bucket and couldn't place the current candidate number into any bucket that would result in our success, we will put it back into the array and return from this DFS path since it was not successful.

            nums.append(candidate)
            return False

Now that our backtracking function is built, all we need to do is pass the initial empty buckets that we are trying to partition.

        return backtrack([0] * k)