LeetCode 698 - Partition to K Equal Sum Subsets

LeetCode Problem 698

Difficulty: 🟡 Medium
Topics: Array, Dynamic Programming, Backtracking, Bit Manipulation, Memoization, Bitmask

Solution

Problem Understanding

The problem is asking whether an array of integers nums can be divided into exactly k subsets such that each subset has the same sum. The input consists of the integer array nums and the integer k. The expected output is a boolean value: true if such a partitioning is possible, and false otherwise.

Key constraints provide important guidance for algorithm design. The array length is at most 16, meaning we can afford exponential-time solutions with bitmasking or backtracking without running into performance issues. Each element can be as large as 10,000, but with only up to 16 elements, the sum remains manageable for integer arithmetic. Each element appears at most 4 times, which restricts the number of duplicate values and helps optimize subset selection.

Important edge cases include arrays with very small lengths, arrays where the sum of elements is not divisible by k (immediately impossible), arrays containing large elements relative to the target sum, and arrays where one element is larger than the sum of the intended subset. These are the inputs most likely to trip up a naive implementation.

Approaches

The brute-force approach involves generating all possible partitions of the array into k subsets and checking whether any of them have equal sums. Conceptually, you could enumerate every possible way to assign each element to a subset, compute the sum for each subset, and return true if all sums match. This guarantees correctness because it exhaustively explores every possibility, but it is computationally infeasible. With n elements and k subsets, there are k^n ways to assign elements to subsets, which is exponential and far too slow for n = 16.

A better approach leverages backtracking with pruning and optionally memoization with bitmasking. The key insight is that if the total sum of all elements is divisible by k, then the target sum for each subset is total_sum / k. We can attempt to build subsets recursively by assigning elements one by one to the current subset, moving to the next subset only when the current one reaches the target sum. Sorting the array in descending order allows early pruning because if the largest element cannot fit in any subset, the partition is impossible. Memoization can store previously seen states of remaining elements to avoid redundant work.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(k^n)	O(n)	Assign each element to a subset and check all combinations. Too slow for n = 16.
Optimal	O(n * 2^n)	O(2^n)	Backtracking with bitmasking and memoization. Explores valid subsets efficiently and prunes impossible paths early.

Algorithm Walkthrough

1. Calculate the total sum of all elements in the array. If the total sum is not divisible by k, immediately return false because equal partitioning is impossible.
2. Determine the target subset sum as target = total_sum / k. This is the sum each subset must achieve.
3. Sort the array in descending order. This ensures larger elements are placed first, which allows early pruning if a number cannot fit into any current subset.
4. Use a bitmask to track which elements are used. Each bit in an integer represents whether a corresponding element is already assigned to a subset. This allows efficient state storage and memoization.
5. Define a recursive function that attempts to fill subsets. The function keeps track of the current bitmask and the sum accumulated for the current subset.
6. Base case: if all elements are used (all bits are 1), return true if the current subset sum is zero. This means all previous subsets have successfully reached the target sum.
7. Iterate through all elements, skipping those already used. Try adding each unused element to the current subset sum. If adding it exceeds the target, skip it. Otherwise, mark the element as used, recurse, and backtrack if the recursion fails.
8. Memoize results for each bitmask to avoid redundant computations.
9. Return the final result of the recursive function.

Why it works: Each recursive call ensures that either the current subset is filled exactly to the target sum or that we explore all possibilities for the remaining elements. The bitmask guarantees that each element is used exactly once, and memoization ensures we do not recompute identical states. Sorting the elements allows pruning impossible paths early, significantly reducing the number of recursive calls.

Python Solution

from typing import List

class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        total_sum = sum(nums)
        if total_sum % k != 0:
            return False
        
        target = total_sum // k
        nums.sort(reverse=True)
        n = len(nums)
        used = 0
        memo = {}
        
        def backtrack(used_mask: int, current_sum: int) -> bool:
            if used_mask == (1 << n) - 1:
                return current_sum == 0
            if (used_mask, current_sum) in memo:
                return memo[(used_mask, current_sum)]
            
            for i in range(n):
                if not (used_mask & (1 << i)) and current_sum + nums[i] <= target:
                    if backtrack(used_mask | (1 << i), (current_sum + nums[i]) % target):
                        memo[(used_mask, current_sum)] = True
                        return True
            
            memo[(used_mask, current_sum)] = False
            return False
        
        return backtrack(0, 0)

The Python solution first checks divisibility and calculates the target sum. Sorting ensures larger numbers are tried first. The recursive backtrack function uses a bitmask to mark used elements and checks all valid combinations. Memoization prevents recomputation. The modulo operation (current_sum + nums[i]) % target allows easy handling of completed subsets.

Go Solution

func canPartitionKSubsets(nums []int, k int) bool {
    total := 0
    for _, num := range nums {
        total += num
    }
    if total%k != 0 {
        return false
    }
    
    target := total / k
    n := len(nums)
    used := 0
    memo := make(map[int]bool)
    
    // Sort nums descending
    for i := 0; i < n; i++ {
        for j := i + 1; j < n; j++ {
            if nums[i] < nums[j] {
                nums[i], nums[j] = nums[j], nums[i]
            }
        }
    }
    
    var backtrack func(int, int) bool
    backtrack = func(usedMask int, currentSum int) bool {
        if usedMask == (1<<n)-1 {
            return currentSum == 0
        }
        key := usedMask*100000 + currentSum
        if val, exists := memo[key]; exists {
            return val
        }
        for i := 0; i < n; i++ {
            if (usedMask&(1<<i)) == 0 && currentSum+nums[i] <= target {
                if backtrack(usedMask|(1<<i), (currentSum+nums[i])%target) {
                    memo[key] = true
                    return true
                }
            }
        }
        memo[key] = false
        return false
    }
    
    return backtrack(0, 0)
}

Go implementation differs slightly because maps cannot directly use tuples as keys, so we combine the usedMask and currentSum into a single integer key. Sorting is done manually since Go has no built-in sort.Sort for reverse integers without importing sort. Otherwise, the logic mirrors the Python version.

Worked Examples

Example 1: nums = [4,3,2,3,5,2,1], k = 4

Target sum = (4+3+2+3+5+2+1)/4 = 20/4 = 5.

Steps:

1. Sort: [5,4,3,3,2,2,1].
2. Place 5 in subset 1: sum = 5 → filled.
3. Place 4 + 1 in subset 2: sum = 5 → filled.
4. Place 3 + 2 in subset 3: sum = 5 → filled.
5. Place remaining 3 + 2 in subset 4: sum = 5 → filled.
6. All subsets filled correctly, return true.

Example 2: nums = [1,2,3,4], k = 3

Total sum = 10, target = 10/3 ≈ 3.33 → not divisible → return false.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n * 2^n)	Each subset state is represented by a bitmask of n elements. We explore valid combinations with memoization.
Space	O(2^n)	Memoization stores results for each bitmask. Bitmask integer and recursion stack also use O(n).

The reasoning is that there are 2^n possible subsets represented by the bitmask. For each subset, we iterate through up to n elements, giving O(n * 2^n) overall.

Test Cases

# Provided examples
assert Solution().canPartitionKSubsets([