Partition to K Equal Sum Subsets: Bucket Backtracking with Pruning

Problem Statement#

Two related problems live under the partition umbrella:

• LeetCode 416 — Partition Equal Subset Sum: given an integer array nums, decide whether it can be split into two subsets with equal sums.
• LeetCode 698 — Partition to K Equal Sum Subsets: given nums and an integer k, decide whether the array can be split into k non-empty subsets each with equal sum.

Both are classic NP-hard problems that ask "can these numbers be balanced into groups?" For small n (up to 16 or so), backtracking with smart pruning crushes them. For larger n, only the 2-bucket case (LC 416) becomes tractable via subset-sum DP.

Why This Problem Matters#

Amazon, Google, and Apple use the partition family to test whether candidates can recognize NP-hardness, choose between backtracking and DP, and apply non-trivial pruning. The 2-subset case (LC 416) has a beautiful O(n times target) DP solution. The k-subset case (LC 698) has no polynomial DP and forces you to write production-grade backtracking with sorting, duplicate skipping, and bucket de-duplication. Candidates who handle both cases on the spot signal that they understand when DP is enough and when backtracking is the only option.

The Core Insight (decision tree / state space)#

For LC 416, the state is (index, remaining_target). The decision is "include nums[index] in subset A or not." That gives a DP table of size O(n times target). Once the table fits in memory, this collapses to a polynomial problem.

For LC 698, the state is more complex: we have k buckets and we are placing one number at a time. The decision is "which bucket does nums[index] go into?" Three pruning ideas keep this tractable:

1. Sort descending — placing large numbers first creates earlier conflicts and prunes faster.
2. Skip oversized adds — if bucket[b] + nums[idx] > target, abandon this branch immediately.
3. De-duplicate equal buckets — at a given recursion node, two empty buckets are interchangeable; trying both wastes time. Use a seen set per call to skip repeats.

These three tricks transform the worst case from k^n to something that finishes within seconds for n up to 16.

Visual Dry Run (recursion tree)#

For nums = [4, 3, 2, 3, 5, 2, 1], k = 4. Sum is 20, target is 5. Sort descending: [5, 4, 3, 3, 2, 2, 1].

place 5: only bucket 0 (others empty -> dedup)
  place 4: bucket 0 full (5+4>5), try bucket 1
    place 3: skip bucket 1 (4+3>5), try bucket 2
      place 3: bucket 2 full, try bucket 3 (empty == bucket 1? no, b1 has 4)
        place 2: bucket 3 has 3, 3+2=5 -> full
          place 2: bucket 1 has 4, 4+2>5; bucket 2 has 3, 3+2=5 -> full
            place 1: bucket 1 has 4, 4+1=5 -> full -> ALL FULL -> true

Notice how each level has very few real choices because of the dedup and oversized-add pruning.

Solution (Optimal) — Python + JavaScript with backtracking template, complexity#

LC 416 (subset-sum DP):

def canPartition(nums):
    total = sum(nums)
    if total % 2:
        return False
    target = total // 2
    dp = {0}
    for x in nums:
        dp |= {s + x for s in dp if s + x <= target}
    return target in dp

LC 698 (bucket backtracking):

def canPartitionKSubsets(nums, k):
    total = sum(nums)
    if total % k:
        return False
    target = total // k
    nums.sort(reverse=True)
    if nums[0] > target:
        return False
    buckets = [0] * k
    n = len(nums)
 
    def backtrack(idx):
        if idx == n:
            return True
        seen = set()
        for b in range(k):
            if buckets[b] in seen:
                continue
            if buckets[b] + nums[idx] > target:
                continue
            seen.add(buckets[b])
            buckets[b] += nums[idx]
            if backtrack(idx + 1):
                return True
            buckets[b] -= nums[idx]
            if buckets[b] == 0:
                break  # if first empty bucket fails, no other empty bucket helps
        return False
 
    return backtrack(0)

function canPartitionKSubsets(nums, k) {
  const total = nums.reduce((a, b) => a + b, 0);
  if (total % k !== 0) return false;
  const target = total / k;
  nums.sort((a, b) => b - a);
  if (nums[0] > target) return false;
  const buckets = new Array(k).fill(0);
  const n = nums.length;
 
  const backtrack = (idx) => {
    if (idx === n) return true;
    const seen = new Set();
    for (let b = 0; b < k; b++) {
      if (seen.has(buckets[b])) continue;
      if (buckets[b] + nums[idx] > target) continue;
      seen.add(buckets[b]);
      buckets[b] += nums[idx];
      if (backtrack(idx + 1)) return true;
      buckets[b] -= nums[idx];
      if (buckets[b] === 0) break;
    }
    return false;
  };
 
  return backtrack(0);
}

Complexity: LC 416 is O(n times target) time, O(target) space. LC 698 is O(k^n) worst case, but the pruning brings it well below that in practice. Space is O(n) recursion plus O(k) buckets.

Common Mistakes#

• Forgetting the if buckets[b] == 0: break shortcut. Without it, every empty bucket gets tried separately.
• Not sorting descending. Ascending sort cripples pruning because small numbers fit everywhere.
• Skipping the total % k != 0 early exit.
• For LC 416, using backtracking when subset-sum DP is asymptotically faster.
• Forgetting nums[0] > target early termination in LC 698.

Interview Tips#

• State both approaches up front: "For 2 subsets I will use subset-sum DP; for k subsets I will use bucket backtracking with three pruning tricks."
• Name the three pruning ideas explicitly: descending sort, oversized-add skip, equal-bucket dedup. Interviewers will press on each.
• Walk through [4,3,2,3,5,2,1], k=4 to show the tree shrinking.
• Mention LeetCode 473 (Matchsticks to Square) — it is just LC 698 with k = 4.

Follow-up Questions#

• Matchsticks to Square (LC 473): same algorithm with k = 4.
• Fair Distribution of Cookies (LC 2305): minimize the maximum bucket sum.
• Subset Sum count: how many subsets reach exactly target? Pure DP.
• What if numbers can be negative? DP space becomes a dict keyed by signed sums.

Key Takeaways#

• Partition problems are NP-hard in general; backtracking with pruning is the practical solution.
• LC 416 (2 subsets) reduces to subset-sum DP — polynomial in target.
• LC 698 (k subsets) needs bucket backtracking with three pruning tricks: descending sort, oversized-add skip, equal-bucket dedup.
• The buckets[b] == 0 break is a non-obvious but critical optimization.
• Mastering this family unlocks Matchsticks to Square, Fair Distribution of Cookies, and the entire bin-packing genre.
• A FAANG interview favorite that tests both pattern recognition and pruning skill.

Sources:

• Partition to K Equal Sum Subsets - LeetCode
• 698. Partition to K Equal Sum Subsets - In-Depth Explanation