Partition to K Equal Sum Subsets (LeetCode 698)

The problem

You get an array nums of positive integers and a number k. Decide whether the numbers can be split into exactly k non-empty groups whose sums are all equal. Every number must be used exactly once. Constraints are small: n ≤ 16.

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

The total is 20, so each group must sum to 5: {5}, {4, 1}, {3, 2}, {3, 2}.

How to think about it

Start with brute force. Assign each number to one of k buckets and check whether all buckets match. That is kⁿ assignments, and most of them are the same partition with the buckets relabelled. For n = 16, k = 4 that is over four billion.

First, cheap pruning. If sum(nums) is not divisible by k, the answer is False. Otherwise every bucket must reach target = sum // k, and any number bigger than target makes it impossible.

The key insight: fill buckets one at a time, and the only thing that matters is which numbers are used. If you always fill the current bucket before opening the next, the state is the set of used numbers plus how much room is left in the current bucket. Two different orders that use the same numbers lead to the same future, so you memoise. With n ≤ 16, the set fits in a bitmask of 2¹⁶ = 65,536 values. When a bucket closes exactly (room hits 0), reset the room to target and start the next bucket.

Which pattern is this? Bitmask DP. The tell is n ≤ 16 together with “use every item exactly once”: the subset of used items is the state.

Approach

1. Compute total. If total % k != 0, return False. Set target = total // k.
2. Sort descending, so big numbers are tried first and failures show up early. If the largest number exceeds target, return False.
3. Define dp(mask, rem): can the unused numbers finish the job, given rem room left in the current bucket?
4. If mask has every bit set, all numbers are placed: return True.
5. Otherwise try each unused number i with nums[i] <= rem. Recurse with bit i set and rem - nums[i], or target if that closes the bucket.
6. Memoise with lru_cache. The answer is dp(0, target).

Solution

from functools import lru_cache
def canPartitionKSubsets(nums, k):
    total = sum(nums)
    if total % k: return False
    target = total // k; n = len(nums); nums.sort(reverse=True)
    if nums[0] > target: return False
    @lru_cache(None)
    def dp(mask, rem):
        if mask == (1 << n) - 1: return True          # every number placed
        for i in range(n):
            if not mask & (1 << i) and nums[i] <= rem:
                nr = rem - nums[i]
                # bucket closed exactly? open a fresh one with full room
                if dp(mask | (1 << i), nr if nr > 0 else target): return True
        return False
    return dp(0, target)
# canPartitionKSubsets([4,3,2,3,5,2,1], 4) -> True

Dry run

After sorting, nums = [5, 4, 3, 3, 2, 2, 1] and target = 5. This is the successful path the recursion takes (indices into the sorted array):

The full mask is reached, so the answer is True. Only 8 states were ever cached on this input.

Complexity

• Time: O(n·2ⁿ). rem is fully determined by the mask (it is target - (sum of used) % target), so there are at most 2ⁿ distinct states, and each tries up to n numbers.
• Space: O(2ⁿ). For the memo table, plus O(n) recursion depth.

Edge cases and pitfalls

• k = 1 is always True; the code handles it naturally because target = total.
• Sum not divisible by k or one number bigger than target: return False before doing any search.
• Duplicates (like the two 3s) cause repeated work in plain backtracking. The memo on mask absorbs most of it.
• Forgetting to reset rem when a bucket closes makes the search think the whole array is one bucket.
• Greedy “biggest first into the emptiest bucket” is wrong. It can paint itself into a corner; you need search or DP.

Related problems

• Partition Equal Subset Sum (LeetCode 416): the k = 2 case, solvable with a 0/1 knapsack instead of a bitmask.
• Shortest Path Visiting All Nodes (LeetCode 847): another “subset of used items” bitmask state.
• Subsets (bitmask enumeration): the basics of treating an integer as a set.