Leetcode 698. Partition to K Equal Sum Subsets

Hadjshell1/3/26About 4 min

Intuition

This is a standard balls-into-bins problem. We make the following assumptions:

Each ball must be placed into exactly one bin.
The bins are unordered (i.e., permutations of bins are considered equivalent).

Given these assumptions, a backtracking approach is appropriate. However, before discussing the implementation, it is important to recognize that there are two equivalent perspectives from which the backtracking can be formulated. Specifically, we can structure the search by:

making choices for each ball (assigning it to a bin or not), or
making choices for each bin (selecting which balls it contains).

These two viewpoints lead to different backtracking implementations. Although they ultimately explore the same solution space and they are mathematically the same, the real codes may have different time complexity.

If you don't know this knowledge, please have a look at my website, where I explain it in detail.

Approach

Before deciding which type of backtracking tree to construct, it is important to consider the following constraint: 1 ≤ k ≤ nums.length ≤ 16.

Since the number of balls can be relatively large, structuring the backtracking from the bin perspective is impractical, as it leads to factorial-time complexity. Instead, we must adopt the ball perspective, which results in an exponential-time backtracking approach and is feasible under the given constraints.

Approach 1

At each decision (chance) node, a ball has k possible choices; that is, each integer in nums can be assigned to one of the k subsets. After all elements in nums have been assigned, if the sums of all buckets are equal, we reach a leaf node in the decision tree and a valid solution is found.

To reduce the size of the decision tree, several pruning strategies can be applied:

Compute the target sum for each subset. If the total sum of nums cannot be evenly divided by k, the function can terminate immediately.
Once a valid solution is found, there is no need to explore the remaining branches of the decision tree.
If a number added into a bucket exceeds the target sum, that bucket should be omitted.
When assigning a ball to a bin, bins with the same current sum are equivalent. Since bins are unordered, it is sufficient to consider only one such bin, thereby reducing redundant choices.

Code

class Solution {
    boolean canPart = false;

    public boolean canPartitionKSubsets(int[] nums, int k) {
        int[] buckets = new int[k];
        int target = target(nums, k);

        if (target == -1)
            return false;

        backtrack(nums, buckets, 0, target);

        return canPart;
    }

    // ball viewpoint
    private void backtrack(int[] nums, int[] buckets, int start_nums, int target) {
        if (canPart)
            return;

        if (start_nums == nums.length) {
            for (int i = 0; i < buckets.length; i++) {
                if (target != buckets[i])
                    return;
            }

            canPart = true;
            return;
        }

        for (int i = 0; i < buckets.length; i++) {
            if (buckets[i] + nums[start_nums] > target)
                continue;

            if (i > 0) {
                boolean isSame = false;

                for (int j = i - 1; j >= 0; j--) {
                    if (buckets[i] == buckets[j]) {
                        isSame = true;
                        break;
                    }
                }

                if (isSame) continue;
            }

            buckets[i] += nums[start_nums];
            backtrack(nums, buckets, start_nums + 1, target);
            buckets[i] -= nums[start_nums];
        }
    }

    private int target(int[] nums, int k) {
        int sum = 0;
        for (int n : nums)
            sum += n;

        return sum % k == 0 ? sum / k : -1;
    }
}

Complexity

Time complexity: $O(k^n)$

Approach 2

Consider the problem of finding a subset of an integer array nums whose sum equals a given integer target. This can be solved via backtracking from the ball's perspective with time complexity $O(2^n)$ .

In the original problem, we can view the task as solving this subset-sum problem $k$ times—once for each subset. This results in an overall time complexity of $O(k*2^n)$ , which is a significant improvement over $O(k^n)$ , especially when $k$ is large.

At each decision (chance) node, a ball has 2 possible choices; that is, each integer in nums can be assigned to the current subset or not. After all buckets have been assigned to the target sum, we reach a leaf node in the decision tree and a valid solution is found. A boolean array used is needed so that the selected number won't be re-assigned into another bucket.

To reduce the size of the decision tree, besides the techniques introduced in Approach 1, another pruning strategy can be applied:

If a subtree of the decision tree—corresponding to a particular bucket configuration—has already been proven to lead to no valid solution, then any equivalent subtree should not be explored again.

For example, let nums = [1, 2, 3, 4, 6] and target = 5. Suppose the first bucket contains [1, 4] and the second bucket contains [2, 3], which does not lead to a valid overall solution. Although the buckets are unordered, naive backtracking treats these assignments as distinct. As a result, it may later explore the symmetric case where the first bucket contains [2, 3] and the second bucket contains [1, 4], which is equally invalid. This exploration is redundant and can be avoided.

To prevent such redundancy, we introduce a hash set to store invalid states (bad choices). If the current bucket configuration has already been recorded as invalid, we stop exploring that branch. Otherwise, we continue the search; if the recursion backtracks without finding a valid solution, we mark the current configuration as a bad choice.

Bitmasking can be used to encode the current selection (e.g., a boolean array used) into an integer, which serves as an efficient hash key for the memoization set.

Code

class Solution {
    boolean canPart = false;

    public boolean canPartitionKSubsets(int[] nums, int k) {
        int[] buckets = new int[k];
        boolean[] used = new boolean[nums.length];
        Set<Integer> badChoices = new HashSet<>();

        int target = target(nums, k);

        if (target == -1)
            return false;

        backtrack(nums, buckets, used, 0, 0, target, badChoices);

        return canPart;
    }

    // ball viewpoint
    private void backtrack(int[] nums, int[] buckets, boolean[] used, int curBucket, int start, int target, Set<Integer> badChoices) {
        if (canPart)
            return;

        if (curBucket == buckets.length) {
            canPart = true;
            return;
        }

        if (buckets[curBucket] == target) {
            int sumTree = hash(used);

            if (!badChoices.contains(sumTree))
                backtrack(nums, buckets, used, curBucket + 1, 0, target, badChoices);

            if (!canPart)
                badChoices.add(sumTree);

            return;
        }

        if (buckets[curBucket] > target || start == nums.length)
            return;

        if (used[start])
            backtrack(nums, buckets, used, curBucket, start + 1, target, badChoices);
        else {
            // not into bucket
            backtrack(nums, buckets, used, curBucket, start + 1, target, badChoices);

            // into bucket
            used[start] = true;
            buckets[curBucket] += nums[start];
            backtrack(nums, buckets, used, curBucket, start + 1, target, badChoices);
            buckets[curBucket] -= nums[start];
            used[start] = false;
        }
    }

    private int hash(boolean[] used) {
        int val = 0;

        for (boolean u : used)
            val = u == true ? val << 1 | 1 : val << 1 | 0;

        return val;
    }

    private int target(int[] nums, int k) {
        int sum = 0;
        for (int n : nums)
            sum += n;

        return sum % k == 0 ? sum / k : -1;
    }
}

Complexity

Time complexity: $O(k*2^n)$