Partition to K Equal Sum Subsets — Bitmask DP for Assignment Problems

Problem Statement#

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.

Constraints:

• 1 <= k <= len(nums) <= 16
• 1 <= nums[i] <= 10^4
• The frequency of each element is in the range [1, 4].

Example 1:

Input:  nums = [4, 3, 2, 3, 5, 2, 1], k = 4
Output: true
Explanation: We can partition into (5), (1, 4), (2, 3), (2, 3) each summing to 5.

Example 2:

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

Example 3:

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

Why This Problem Matters#

Partition to K Equal Sum Subsets is the canonical bitmask DP problem: it asks you to make a series of choices over a subset of elements, and the DP state must track exactly which elements have been chosen. With n up to 16, there are 2^n = 65,536 possible subsets — small enough to enumerate as bitmasks, large enough to make naive DFS too slow without memoization.

Why FAANG interviewers love this problem: it reveals whether a candidate can compress a multi-dimensional state (which subset is filled, plus current bucket fill level) into a single bitmask. Once you see this trick, harder problems like Smallest Sufficient Team (LC 1125), Stickers to Spell Word (LC 691), and Traveling Salesman become approachable.

The problem also has a deep connection to bin packing and multiprocessor scheduling, both NP-hard. The bitmask DP gives a pseudo-polynomial algorithm in n that is competitive with sophisticated branch-and-bound approaches for small n.

The Core Insight#

Key observation: the order in which we fill the k buckets does not matter. So we can fill them one at a time. We greedily fill bucket 1 to capacity target = total / k, then bucket 2, then bucket 3, and so on.

State: dp[mask] = the running sum of the currently filling bucket when the set of used elements is exactly mask. We use -1 to mean "this state is unreachable."

Transition: From state dp[mask], try adding each unused element nums[i]:

• New state: mask | (1 << i).
• New running sum: (dp[mask] + nums[i]) % target. The % target does the magic — when the bucket is exactly filled, the value resets to 0, which means we have started a new (empty) bucket.
• The transition is valid only if dp[mask] + nums[i] <= target (we never overflow a bucket).

Answer: dp[(1 << n) - 1] == 0 means all elements are used and the last bucket was filled exactly. If this final state is reachable, the answer is true.

This converts an exponential DFS (each element has k placements) into a clean DP over 2^n masks with n transitions per mask.

Visual Dry Run#

Input: nums = [4, 3, 2, 3, 5, 2, 1], k = 4. Sum = 20, target = 5.

After sorting descending: [5, 4, 3, 3, 2, 2, 1].

Initial state: dp[0000000] = 0 (no elements used, current bucket has sum 0).

Selected transitions (showing only successful ones):

• From 0000000 (sum 0), pick 5 at index 0: dp[1000000] = (0 + 5) % 5 = 0. Bucket filled, restart.
• From 1000000 (sum 0), pick 4 at index 1: dp[1100000] = (0 + 4) % 5 = 4.
• From 1100000 (sum 4), pick 1 at index 6: dp[1100001] = (4 + 1) % 5 = 0. Bucket filled.
• From 1100001 (sum 0), pick 3 at index 2: dp[1110001] = 3.
• From 1110001, pick 2 at index 4: dp[1111001] = 5 % 5 = 0. Filled.
• From 1111001, pick 3 at index 3: dp[1111101] = 3.
• From 1111101, pick 2 at index 5: dp[1111111] = 5 % 5 = 0. Filled.

Final mask 1111111 reached with dp = 0 — all elements used, last bucket exact. Return true.

Solution (Optimal)#

Python#

class Solution:
    def canPartitionKSubsets(self, nums: list[int], k: int) -> bool:
        total = sum(nums)
        # Quick reject: total must split evenly into k buckets
        if total % k:
            return False
        target = total // k
 
        # Sort descending so we hit big elements first; better pruning
        nums.sort(reverse=True)
        if nums[0] > target:
            return False
 
        n = len(nums)
        # dp[mask] = sum in the currently-filling bucket; -1 if unreachable
        dp = [-1] * (1 << n)
        dp[0] = 0
 
        for mask in range(1 << n):
            if dp[mask] == -1:
                continue
            for i in range(n):
                # skip if element i is already used in this mask
                if mask & (1 << i):
                    continue
                # adding nums[i] must not overflow the bucket
                if dp[mask] + nums[i] <= target:
                    next_mask = mask | (1 << i)
                    # mod target rolls the bucket over to a fresh one
                    dp[next_mask] = (dp[mask] + nums[i]) % target
 
        # All elements used and last bucket exactly filled
        return dp[(1 << n) - 1] == 0

JavaScript#

var canPartitionKSubsets = function(nums, k) {
    const total = nums.reduce((a, b) => a + b, 0);
    // Total must divide evenly into k equal buckets
    if (total % k !== 0) return false;
    const target = total / k;
 
    // Sort descending for stronger pruning
    nums.sort((a, b) => b - a);
    if (nums[0] > target) return false;
 
    const n = nums.length;
    // dp[mask] = sum in the bucket currently being filled
    const dp = new Array(1 << n).fill(-1);
    dp[0] = 0;
 
    for (let mask = 0; mask < (1 << n); mask++) {
        if (dp[mask] === -1) continue;
        for (let i = 0; i < n; i++) {
            if (mask & (1 << i)) continue;          // element already used
            if (dp[mask] + nums[i] > target) continue; // would overflow bucket
            const next = mask | (1 << i);
            // mod target resets the bucket when it fills exactly
            dp[next] = (dp[mask] + nums[i]) % target;
        }
    }
 
    return dp[(1 << n) - 1] === 0;
};

Complexity: Time O(2^n * n), Space O(2^n).

Common Mistakes#

1. Skipping the divisibility check. If total % k != 0, no partition exists. Forgetting this check wastes work and can produce wrong answers if you test a target that does not exist.

2. Forgetting the nums[0] > target early-exit. After sorting descending, if the largest element exceeds target, no bucket can hold it. Skipping this prune still works but slows the DP dramatically.

3. Using dp[mask] + nums[i] < target instead of <= target. The strict inequality misses the case when an element exactly fills a bucket ((dp + nums[i]) % target == 0).

4. Not using % target. Without the modulo, dp[mask] becomes the total of all elements used, not the running bucket sum. The DP no longer tracks bucket boundaries.

5. Sorting ascending instead of descending. Descending sort hits big elements first, pruning impossible branches earlier. Ascending sort still works but loses orders of magnitude in practice.

6. Confusing this with Partition Equal Subset Sum (LC 416). LC 416 is k = 2 and uses 1D subset-sum DP. LC 698 generalizes to arbitrary k and requires bitmask state.

Interview Tips#

• Frame the problem as "fill buckets one at a time." This greatly simplifies the explanation versus thinking of all k buckets in parallel.
• Explain why ordering doesn't matter: any partition is just a relabeling of buckets, so we can WLOG fill them sequentially.
• Mention the alternative DFS-with-memoization solution. It is equivalent in complexity but more verbose; bitmask DP shines for its conciseness.
• Discuss pruning: descending sort, early nums[0] > target rejection, and skipping any mask with dp[mask] == -1.
• For larger n (16 to 32), mention meet-in-the-middle or backtracking with strong pruning as alternatives.

Follow-up Questions#

Q: What if n could be 30+? 2^30 = 1 billion, too large for the DP table. Use backtracking with aggressive pruning (skip duplicate elements, pre-sort, abandon branches when bucket overshoots).

Q: How does this relate to bin packing? This problem is the decision version of multiway partition / bin packing where you must use exactly k bins of capacity target. Bitmask DP gives an O(2^n * n) algorithm, far better than naive O(k^n).

Q: How would you reconstruct one valid partition? Store a parent pointer for each dp[mask]. After confirming dp[(1 << n) - 1] == 0, walk back through parents to identify which element each bucket received.

Q: Why is dp[mask] = (prev + nums[i]) % target correct? When the bucket would be exactly filled, prev + nums[i] = target, so mod target = 0 — which represents starting a new empty bucket. When not yet filled, mod target is a no-op since prev + nums[i] < target.

Key Takeaways#

• Bitmask DP encodes "which elements are used" in a single integer mask, perfect for n up to about 20.
• dp[mask] stores the running sum of the currently filling bucket, with the elegant % target trick to roll over to a new bucket.
• Sort descending and prune nums[0] > target before running the DP for major speedups.
• Final answer is dp[(1 << n) - 1] == 0: all elements used and last bucket exactly full.
• The pattern generalizes to assignment, bin packing, and multiprocessor scheduling problems.
• Time O(2^n * n), space O(2^n) — optimal for this state-space size.