Partition to K Equal Sum Subsets

Medium Top 250

Associated Patterns

Interviewed At (Company Tags)

AmazonGoogleMicrosoft

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.

Examples

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).

Example 2:

Input: nums = [1,2,3,4], k = 3 Output: false

Constraints

1 ≤ k ≤ nums.length ≤ 16
1 ≤ nums[i] ≤ 10⁴
The frequency of each element is in the range [1, 4]

Fill Each Bucket to the Target

The total must be divisible by k. Each subset should sum to total / k. The backtracking tries to fill k buckets one at a time, placing each unused number into the current bucket, and moving to the next bucket once the current one reaches the target.

The key pruning: if you can’t fill the first bucket, you can’t fill any. Sort descending so large numbers are tried first and dead branches fail early.

A bitmask memoization on the used state can reduce repeated work: if the same set of remaining elements was tried and failed before, skip it. This brings the worst case down significantly.

Solution 1: Backtracking with Pruning

import java.util.*;

class Solution {
    public boolean canPartitionKSubsets(int[] nums, int k) {
        int total = 0;
        for (int n : nums) total += n;
        if (total % k != 0) return false;

        int target = total / k;
        Arrays.sort(nums);
        if (nums[nums.length - 1] > target) return false; // largest element exceeds target

        // sort descending for better pruning
        int[] sorted = new int[nums.length];
        for (int i = 0; i < nums.length; i++) sorted[i] = nums[nums.length - 1 - i];

        return backtrack(sorted, new boolean[sorted.length], k, 0, 0, target);
    }

    private boolean backtrack(int[] nums, boolean[] used, int k,
                                int index, int currentSum, int target) {
        if (k == 0) return true; // all buckets filled
        if (currentSum == target) return backtrack(nums, used, k - 1, 0, 0, target);

        for (int i = index; i < nums.length; i++) {
            if (used[i]) continue;
            if (currentSum + nums[i] > target) continue;
            // skip duplicates at the same bucket level
            if (i > index && !used[i - 1] && nums[i] == nums[i - 1]) continue;

            used[i] = true;
            if (backtrack(nums, used, k, i + 1, currentSum + nums[i], target)) return true;
            used[i] = false;
        }
        return false;
    }
}

class Solution:
    def canPartitionKSubsets(self, nums: list[int], k: int) -> bool:
        total = sum(nums)
        if total % k != 0:
            return False

        target = total // k
        nums.sort(reverse=True)

        if nums[0] > target:
            return False

        used = [False] * len(nums)

        def backtrack(k: int, index: int, current_sum: int) -> bool:
            if k == 0:
                return True
            if current_sum == target:
                return backtrack(k - 1, 0, 0)

            for i in range(index, len(nums)):
                if used[i]:
                    continue
                if current_sum + nums[i] > target:
                    continue
                if i > index and not used[i - 1] and nums[i] == nums[i - 1]:
                    continue  # skip duplicate

                used[i] = True
                if backtrack(k, i + 1, current_sum + nums[i]):
                    return True
                used[i] = False

            return False

        return backtrack(k, 0, 0)

#include <vector>
#include <algorithm>
#include <numeric>

class Solution {
public:
    bool canPartitionKSubsets(std::vector<int>& nums, int k) {
        int total = std::accumulate(nums.begin(), nums.end(), 0);
        if (total % k != 0) return false;

        int target = total / k;
        std::sort(nums.rbegin(), nums.rend()); // descending

        if (nums[0] > target) return false;

        std::vector<bool> used(nums.size(), false);
        return backtrack(nums, used, k, 0, 0, target);
    }

private:
    bool backtrack(std::vector<int>& nums, std::vector<bool>& used, int k,
                   int index, int currentSum, int target) {
        if (k == 0) return true;
        if (currentSum == target) return backtrack(nums, used, k - 1, 0, 0, target);

        for (int i = index; i < (int)nums.size(); i++) {
            if (used[i]) continue;
            if (currentSum + nums[i] > target) continue;
            if (i > index && !used[i - 1] && nums[i] == nums[i - 1]) continue;

            used[i] = true;
            if (backtrack(nums, used, k, i + 1, currentSum + nums[i], target)) return true;
            used[i] = false;
        }
        return false;
    }
};

Complexity Analysis:

Time Complexity: O(k * 2^n) in the worst case. With pruning, practical performance is much better.
Space Complexity: O(n) for the used array and recursion stack.

Where it breaks: the duplicate-skip prune nums[i] == nums[i-1] && !used[i-1] is the same pattern as Permutations II/Subsets II. Without it, redundant recursive calls on equivalent elements multiply the runtime significantly on inputs with repeated values.

Common Mistakes

Checking total % k != 0 but forgetting nums[0] > target after sorting descending. A single element larger than the target makes the partition impossible.
Resetting index to 0 when filling a new bucket. This is intentional: a new bucket can include any unused element from the beginning. Continuing from index (the old position) incorrectly prevents earlier elements from joining a new bucket.

Frequently Asked Questions

Why is this problem NP-complete? Partition to K Equal Sum Subsets is a generalization of the subset sum problem, which is NP-complete. Backtracking with pruning is the standard practical approach; there’s no known polynomial-time algorithm.

How does this relate to Matchsticks to Square? Matchsticks to Square is this problem with k = 4. The algorithm is identical.