Skip to content

698. Partition to K Equal Sum Subsets 👍

  • Time: $O(2^n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution {
 public:
  bool canPartitionKSubsets(vector<int>& nums, int k) {
    const int sum = accumulate(nums.begin(), nums.end(), 0);
    if (sum % k != 0)
      return false;

    const int target = sum / k;  // the target sum of each subset
    if (ranges::any_of(nums, [target](const int num) { return num > target; }))
      return false;

    ranges::sort(nums, greater<>());
    return dfs(nums, 0, k, /*currSum=*/0, target, /*used=*/0);
  }

 private:
  bool dfs(const vector<int>& nums, int s, int remainingGroups, int currSum,
           const int subsetTargetSum, int used) {
    if (remainingGroups == 0)
      return true;
    if (currSum > subsetTargetSum)
      return false;
    if (currSum == subsetTargetSum)  // Find a valid group, so fresh start.
      return dfs(nums, 0, remainingGroups - 1, 0, subsetTargetSum, used);

    for (int i = s; i < nums.size(); ++i) {
      if (used >> i & 1)
        continue;
      if (dfs(nums, i + 1, remainingGroups, currSum + nums[i], subsetTargetSum,
              used | 1 << i))
        return true;
    }

    return false;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class Solution {
  public boolean canPartitionKSubsets(int[] nums, int k) {
    final int sum = Arrays.stream(nums).sum();
    if (sum % k != 0)
      return false;

    final int target = sum / k; // the target sum of each subset
    if (Arrays.stream(nums).anyMatch(num -> num > target))
      return false;

    return dfs(nums, 0, k, /*currSum=*/0, target, /*used=*/0);
  }

  private boolean dfs(int[] nums, int s, int remainingGroups, int currSum, int subsetTargetSum,
                      int used) {
    if (remainingGroups == 0)
      return true;
    if (currSum > subsetTargetSum)
      return false;
    if (currSum == subsetTargetSum) // Find a valid group, so fresh start.
      return dfs(nums, 0, remainingGroups - 1, 0, subsetTargetSum, used);

    for (int i = s; i < nums.length; ++i) {
      if ((used >> i & 1) == 1)
        continue;
      if (dfs(nums, i + 1, remainingGroups, currSum + nums[i], subsetTargetSum, used | 1 << i))
        return true;
    }

    return false;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution:
  def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
    summ = sum(nums)
    if summ % k != 0:
      return False

    target = summ // k  # the target sum of each subset
    if any(num > target for num in nums):
      return False

    def dfs(s: int, remainingGroups: int, currSum: int, used: int) -> bool:
      if remainingGroups == 0:
        return True
      if currSum > target:
        return False
      if currSum == target:  # Find a valid group, so fresh start.
        return dfs(0, remainingGroups - 1, 0, used)

      for i in range(s, len(nums)):
        if used >> i & 1:
          continue
        if dfs(i + 1, remainingGroups, currSum + nums[i], used | 1 << i):
          return True

      return False

    nums.sort(reverse=True)
    return dfs(0, k, 0, 0)