Partition to K Equal Sum Subsets: Backtracking with Pruning

Given an integer array nums and an integer k, return true if you can divide the array into k non-empty subsets whose sums are all equal. Otherwise, return false.

This is LeetCode 698: Partition to K Equal Sum Subsets, and it is one of the best problems for learning how to combine backtracking with aggressive pruning. You are not just generating subsets here. You are distributing elements across multiple buckets under tight constraints, and the naive approach is far too slow without the right optimizations.

Why this problem matters

Partition to K Equal Sum Subsets sits at a critical intersection of backtracking, pruning, and optimization. If you have already solved Subsets and Combination Sum, you know how to generate combinations and explore decision trees. This problem forces you to take those skills further.

The core challenge is not conceptual. The idea is simple: split the numbers into k groups with equal sums. The challenge is computational. A brute-force search over all possible partitions is exponential, and without smart pruning, even moderate inputs will time out. This makes it an excellent problem for learning how the right pruning strategies can take an infeasible solution and make it fast enough to pass.

This problem also connects directly to Matchsticks to Square, which is the special case where k = 4. If you can solve Partition to K Equal Sum Subsets, Matchsticks to Square becomes trivial. More broadly, the bucket-filling backtracking pattern shows up in scheduling problems, resource allocation, and any situation where you need to distribute items across groups with constraints.

The approach: backtracking with buckets

The first thing to check: does a valid partition even exist? If the total sum of the array is not divisible by k, the answer is immediately false. Otherwise, the target sum for each bucket is total / k.

Now picture k empty buckets. You iterate through each number and try placing it into one of the k buckets. If a bucket's running sum would exceed the target after adding the number, skip that bucket. If you successfully place all numbers such that every bucket hits the target, return true. If you exhaust all options for a number without finding a valid placement, backtrack.

def canPartitionKSubsets(nums, k):
    total = sum(nums)
    if total % k != 0:
        return False

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

    if nums[0] > target:
        return False

    buckets = [0] * k

    def backtrack(index):
        if index == len(nums):
            return all(b == target for b in buckets)

        for i in range(k):
            if buckets[i] + nums[index] <= target:
                buckets[i] += nums[index]
                if backtrack(index + 1):
                    return True
                buckets[i] -= nums[index]

            if buckets[i] == 0:
                break

        return False

    return backtrack(0)

This solution sorts the array in descending order so that larger numbers get placed first. Larger numbers constrain the search more aggressively, which means invalid branches get pruned earlier. The if buckets[i] == 0: break line is a key optimization: if placing a number into an empty bucket did not lead to a solution, there is no point trying it in a different empty bucket, because all empty buckets are interchangeable at that point.

Optimizing with pruning

The basic backtracking solution above already includes two important pruning strategies, but let's make them explicit and add a few more.

1. Sort descending. Large numbers are harder to fit into buckets. Placing them first means you discover dead ends sooner. If the largest number exceeds the target, you can return false immediately without any recursion.

2. Skip duplicate bucket states. When multiple buckets have the same current sum, placing a number into any one of them is equivalent. The if buckets[i] == 0: break line handles the most common case (all empty buckets are identical), but you can extend this: if buckets[i] == buckets[j] for i < j, skip bucket j.

3. Early termination on filled buckets. When a bucket's sum equals the target, it is full. You never need to add anything else to it. The buckets[i] + nums[index] <= target check handles this implicitly.

4. Fail fast on impossible elements. If any single element is larger than the target, no valid partition exists. Check this before starting the recursion.

Here is the optimized version with all pruning in place:

def canPartitionKSubsets(nums, k):
    total = sum(nums)
    if total % k != 0:
        return False

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

    if nums[0] > target:
        return False

    buckets = [0] * k

    def backtrack(index):
        if index == len(nums):
            return True

        seen = set()
        for i in range(k):
            if buckets[i] in seen:
                continue
            if buckets[i] + nums[index] <= target:
                seen.add(buckets[i])
                buckets[i] += nums[index]
                if backtrack(index + 1):
                    return True
                buckets[i] -= nums[index]

        return False

    return backtrack(0)

The seen set tracks bucket sums we have already tried for the current number. If two buckets have the same sum, placing the number in either one would produce identical future search trees, so we skip the duplicate. This single optimization can dramatically reduce the number of branches explored.

Notice that the base case simplified too. When index == len(nums), all numbers have been placed. Since we only ever place a number when buckets[i] + nums[index] <= target, and the total sum equals k * target, reaching the end means every bucket must be exactly at the target.

Complexity analysis

Where n is the number of elements in nums.

The backtracking approach has a worst-case time of O(k^n) because each of the n elements can go into any of the k buckets. In practice, the pruning strategies reduce this dramatically. Sorting descending and skipping duplicate bucket states often bring the runtime down to something very manageable for the LeetCode constraints (n <= 16).

The bitmask DP approach uses a bitmask to represent which elements have been used. For each bitmask, it tracks whether that subset of elements can be partitioned into complete buckets. This gives a cleaner O(n * 2^n) time bound, but uses O(2^n) space. For n = 16, that is 65,536 states, which is fine. The bitmask DP is theoretically more predictable in runtime, but the pruned backtracking approach is often faster in practice for typical inputs.

Edge cases to watch for

• Total not divisible by k: return false immediately. No partition is possible.
• Any element larger than target: return false. That element cannot fit in any bucket.
• k = 1: always true (the entire array is one subset).
• k = n: only true if every element is equal.
• All elements the same: true if n is divisible by k, since each bucket gets the same count of identical elements.
• Single element array: only true if k = 1.
• Large values: the elements can be up to 10,000, but since n <= 16 the total is bounded. The target can be large, but bucket sums are just integers.

The building blocks

This problem is built on the choose-explore-unchoose backtracking template, the same pattern you use in Subsets, Combination Sum, and Word Search. The skeleton is always the same:

def backtrack(state):
    if is_complete(state):
        return is_valid(state)

    for choice in candidates(state):
        make_choice(choice)
        if backtrack(state):
            return True
        undo_choice(choice)

    return False

In Partition to K Equal Sum Subsets, make_choice is adding a number to a bucket. undo_choice is removing it. candidates are the buckets that can accept the number without exceeding the target. The twist compared to simpler backtracking problems is that you are distributing elements across multiple containers instead of building a single subset, and the pruning strategies are what make the difference between a solution that passes and one that times out.

The pruning itself follows a general principle that applies to many backtracking problems: skip equivalent states. Whenever two branches of the search tree will produce the same outcome, you only need to explore one. In this problem, that means skipping duplicate bucket sums. In Subsets II, it means skipping consecutive duplicate elements. The strategy is the same, just applied to different types of symmetry.

From understanding to recall

You have seen the bucket-filling approach, the pruning strategies, and why each one matters. But in an interview, you need to produce this from memory. The tricky parts are not the high-level idea. They are the details: remembering to sort descending, implementing the duplicate bucket skip, getting the base case right.

Spaced repetition bridges the gap between understanding and fluent recall. You practice writing the backtracking skeleton, then layer on the pruning optimizations, and after a few rounds the whole solution flows out naturally. When someone says "partition into k subsets," you do not hesitate. You know exactly how to set up the buckets, where to prune, and why.

The choose-explore-unchoose pattern is one of roughly 60 reusable building blocks that cover hundreds of LeetCode problems. Drilling them individually with spaced repetition beats random problem grinding every time.

Related posts

• Subsets - The foundation for all backtracking subset problems, using the same choose-explore-unchoose template
• Combination Sum - Backtracking with a target constraint, same pattern with element reuse allowed
• Target Sum - Another partition-style problem where you distribute elements into two groups, solvable with both backtracking and DP