Partition to k equal sum subsets

Jhanak Didwania
TRICK THE INTERVIEWER
4 min readMay 24, 2020

Hi guys, I am back with another coding interview problem. So, here we go!

Problem statement: Given an array of integers nums and a positive integer k, find whether it's possible to divide this array into k non-empty subsets whose sums are all equal.

This question is an extension of the ‘Equal subset-sum’ problem where the goal was to find whether it is possible to divide an array into two non-empty subsets whose sums are equal. In that problem, k is fixed as 2. So that question basically becomes a classical subset-sum problem where the goal is to find whether there exists a subset of the given array whose sum is sum/2.

The question is NP-hard as there are total 2^n subsets where n is the size of the array.

Partition to k equal sum subsets

nums, size = 7 and #of partitions, k = 4

First check whether it is possible to make k subsets of the array.

  1. If the length of the array is less than the partitions required, then we can’t partition it and return false.
    if(nums.length<k) return false;
  2. Secondly, find the sum required for “k” individual buckets. If this quantity is not an integer, the task is impossible. Initialize newSum as zero. This newSum will denote the sum for each bucket.
    for(int i=0; i<nums.length; i++){
    newSum+=nums[i];
    }
    if(newSum%k !=0) return false;
    newSum = newSum/k;

newSum : 20
newSum = 20/4 = 5

Therefore, in this question we need to find four non-empty subsets such that the sum of each subset is 5. These are our 4 subsets: {4,1}, {5}, {3,2}, {3,2}

Backtracking approach

We have k buckets. In this example say, 4. For each element of the array, we have two choices, either we include the element in the current bucket, or we can exclude it from the current bucket.

We can only include the element in the current bucket if:

  1. Element is unused
  2. currentElement + currentSum ≤ target bucket sum

Suppose, we introduce a new boolean array, it will keep track of the elements that are used so far, so that, we do not reuse the same element again.

Initialize all the elements of the used array to false as no element has been used so far

We define a recursive function, Partition that will return whether it’s possible to partition the given array into k subsets such that the sum of all is equal.

We pass 6 parameters in this function.
k = number of remaining buckets to be filled
sum = newSum (total sum required in each bucket)
cs = current sum of the current bucket (initially zero)
nums = array of integers whose subsets are to be found
used = to keep track of how many elements have been used so far
start = starting index of the array (initially zero)

private boolean Partition(int k, int sum, int cs, int []nums, boolean used[], int start){

/** base condition, only one bucket remains, therefore no need to recurse further **/
if(k==1) return true;

/** if the sum of elements of current bucket is equal to required sum then current bucket is considered to be filled, and then we try to fill the next bucket. In this case, we reduce number of remaining buckets by one and we initialise current sum for next bucket to be zero **/
if(cs == sum) return Partition(k-1, sum, 0, nums, used, 0);

/** Loop for filling the current bucket **/
for(int i=start; i<nums.length; i++){
if(!used[i] && cs+nums[i]<=sum){
used[i] = true; //include the current unused element and mark it as used
/** recurse for the remaining array to fill the current bucket, bucket sum will increase to cs+nums[i] including the current elements and start index for next iteration will be current index+1 **/ if(Partition(k, sum, cs+nums[i], nums, used, i+1)) return true; used[i] = false; //exclude the current element and mark it as unused, backtracking. We will do backtracking if on considering the current item in the bucket, we can’t fill the current bucket or the remaining buckets. Then we move on adding the next element in the current bucket and remove this one.
}
}

return false; //if partitions are not found, return false
}

Find below a dry run of the algorithm for the example array:

Notice, cs has been set to zero again and k has been reduced to 3 as we are decreasing the remaining buckets count, we will continue this process till number of remaining bucket equals one

Similarly, we can recur for the remaining buckets.

Space Complexity: O(N), the space used by recursive calls to search in our call stack and O(N) for the boolean array used to keep track of elements.

I hope the algorithm now makes sense to you!. If you like this blog, please share it and your claps always motivate me. Happy Coding! :D

--

--