30/09/2021

[Java] Partition array in K subsets of equal sum

Given an array of positive integers and a target K, determine whether it's possible to partition the array in K subsets of equal sum.

Example:

1,1,1,1 k = 2

it's possible to have two sets, each with a total sum of 2

1,1,1,1 k = 3

it's not possible to have 3 sets of equal sum

We first determine if this is generally feasible by looking at the total sum of elements in our array and comparing with K. If possible, each set should sum up to total / K, which means K must be a divisor of our total sum.

We then need to explore all possibilities to generate each subset, this is evident in this sample:

1,1,1,2,2,2 k = 3

The only solution is to pair each 1 with a 2, we cannot apply a greedy approach as we might miss this solution.

A bit similar to the house painting problem, we need to recursively explore our options keeping track of:

  • number of "full" subsets, that is subsets we created so far that reached our target sum
  • available elements, as we can only place each element in a single subset
  • current sum of each subset

We start with no subset existing and all elements available and try an add each available element to a set until we reach the target sum. Everytime we use an element, we mark it as unavailable.

When a set reaches the target sum, we increase the number of "full" subsets and start creating a new subset repeating this logic.

When we have no more available elements, if the amount of "full" subsets is reached, we have our answer, otherwise we need to backtrack and explore a different configuration.

When we backtrack we simply mark the current element as available again, remove it from the current subset sum and proceed on the next element.

This approach gives us a (N * 2^N) time complexity as for each element we explore the option of choosing it or not for the current subset and worst case we repeat this process for each element in our array before we realize no solution exists. Not sure whether sorting the array would have a significant impact on the runtime since there can exist a configuration where the total sum is divisible by K but no arrangement of elements is possible, maybe sorting descending can on average be faster but not in worst case.

We use O(N) space instead for the recursion stack (can't have more than N calls as we pick each element once in a branch) and the array tracking the available elements.

You can check my implementation of canPartitionInKSubsetsOfEqualSum on my Gist along with some tests in CanPartitionInKSubsetsOfEqualSumJTests.

No comments:

Post a Comment

With great power comes great responsibility