Skip to main content

Partition to K Equal Sum Subsets

Partition to K Equal Sum Subsets: 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 is 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].

Try this Problem on your own or check similar problems:

  1. Partition Equal Subset Sum
  2. Fair Distribution of Cookies
  3. Maximum Number of Ways to Partition an Array
Solution:​
public boolean canPartitionKSubsets(int[] nums, int k) {
    int n = nums.length, sum = 0;
    int[] dp = new int[1<<16];
    Arrays.fill(dp, -1);
    dp[0] = 0;

    for (int i = 0; i < n; i++) sum += nums[i];
    if (sum % k != 0) return false;

    int target = sum / k;

    for (int mask = 0; mask < (1<<n); ++mask) {
        if (dp[mask] == -1) continue;
        for (int i = 0; i < n; ++i) {
            if ((mask & ( 1<<i )) == 0 && dp[mask]+nums[i] <= target) {
                dp[mask|(1<<i)] = (dp[mask] + nums[i]) % target;
            }
        }
    }
    return dp[(1<<n)-1] == 0;
}
Time/Space Complexity:​
  • Time Complexity: O(n2^n)
  • Space Complexity: O(1<<16)
Explanation:​

Let's first start with the backtracking solution:

class Solution {
  public boolean canPartitionKSubsets(int[] nums, int k) {
      int sum = 0;
      for (int i = 0; i < nums.length; ++i) {
          sum += nums[i];
      }
      if(sum % k != 0)  return false;
      Arrays.sort(nums);
      return helper(nums, nums.length - 1, 0, k, 0, sum / k, new boolean[nums.length]);
  }

  private boolean helper(int[] nums, int start, int round, int k, int currentSum, int targetSum, boolean[] visited){
      if(round == k) return true;
      if(currentSum == targetSum) return helper(nums, nums.length - 1, round + 1, k, 0, targetSum, visited);

      for(int i = start; i >= 0; --i){
          if(!visited[i] && currentSum + nums[i] <= targetSum){
              visited[i] = true;
              if(helper(nums, i - 1,round, k, currentSum + nums[i], targetSum, visited)) return true;
              visited[i] = false;
          }
      }
      return false;
  }
}

The algorithm above works as follows:

  • We sum all the numbers from the input array and check if the sum can be equally divided into k subsets sum % k == 0
  • We sort the array to start from the largest number first (since there are many more ways to combine smaller numbers into subset so the sum is equal to k than there is when we're using bigger number)
  • Next we run the helper function recursively as always and check if we can group the numbers to form subsets of sums sum/k
  • It's pretty much similar if we would search for one sum instead of k sums, but this time we use round to mark how many subsets we have already found with the sum sum/k, if the round is k than we can break the input array into k subsets
  • We use array visited to track which numbers we have already used, and we only start the next round if we found a subset with sum/k

Now let's move to the final DP solution and break it down:

  • We do the same checking if the total sum is divisible by k
  • By the problem statement we know that there max 16 elements so we can create dp where each number represent if one of the elements in the input array is chosen or not (e.g nums = [1,2,3], dp["011"] means that 1 is not present, while 2 and 3 are)
  • Note this is the same as finding subsets, there are 2^n (or 1<<n) of them for input array of length n, and this is the total number of our masks
  • With dp[mask] == -1 we only check valid masks that we get by choosing one of the input arrays elements, so in the first iteration we will be looking at what happens if we choose each of the elements one by one
  • mask & ( 1<<i )) == 0 this statement checks if have already chosen the element, and if not, we only choose it if the current value for the mask before selecting this element when added to the value of that element is smaller or equal to target (so we don't go over the target for one of subsets sums)
  • If the requirements are fulfilled, we set the value of new mask (with the new element selected as well) to (dp[mask] + nums[i]) % target, we are only interested in remainder since the whole parts will form subsets of sums equal to target == sum / k
  • At the end we check if the value of dp["1111...11"] (all elements selected) is equal to 0 (no remainder, array can be divided in subsets of sum sum/k)
public boolean canPartitionKSubsets(int[] nums, int k) {
    int n = nums.length, sum = 0;
    int[] dp = new int[1<<16];
    Arrays.fill(dp, -1);
    dp[0] = 0;

    for (int i = 0; i < n; i++) sum += nums[i];
    if (sum % k != 0) return false;

    int target = sum / k;

    for (int mask = 0; mask < (1<<n); ++mask) {
        if (dp[mask] == -1) continue;
        for (int i = 0; i < n; ++i) {
            if ((mask & ( 1<<i )) == 0 && dp[mask]+nums[i] <= target) {
                dp[mask|(1<<i)] = (dp[mask] + nums[i]) % target;
            }
        }
    }
    return dp[(1<<n)-1] == 0;
}