BackTrackingProblem 18 of 19

Partition of a set into K subsets with equal sum

Brute Force
Time: O(k^n)
Space: O(n)
Optimal
Time: O(k × 2^n)
Space: O(k × n)

Problem Statement

Given an integer array of n elements and an integer k, partition the array into k subsets of equal sum. Return true if such partition is possible, otherwise return false.

Example:

Input: arr = [4, 3, 2, 3, 5, 2, 1], k = 4 Output: true Explanation: Subsets: [5], [1, 4], [2, 3], [2, 3] Each subset sums to 5. Input: arr = [1, 2, 3, 4], k = 3 Output: false Explanation: Cannot divide into 3 equal sum subsets.

Approach 1: Brute Force (Assign Each Element to Bucket)

Intuition

Try assigning each element to one of k buckets. If all buckets reach the target sum exactly, return true. Use backtracking when assignment fails.

Algorithm

  1. Calculate target sum = totalSum / k
  2. Try assigning each element to each bucket
  3. If bucket + element <= target, add and recurse
  4. When all elements assigned, check if all buckets equal target
  5. Backtrack if assignment fails
java
import java.util.*;

class Solution {
    public boolean canPartition(int[] arr, int index, int[] buckets, int k, int target) {
        // All elements assigned
        if (index == arr.length) {
            for (int i = 0; i < k; i++) {
                if (buckets[i] != target) return false;
            }
            return true;
        }
        
        // Try adding current element to each bucket
        for (int i = 0; i < k; i++) {
            if (buckets[i] + arr[index] <= target) {
                buckets[i] += arr[index];
                
                if (canPartition(arr, index + 1, buckets, k, target)) {
                    return true;
                }
                
                buckets[i] -= arr[index];  // Backtrack
            }
            
            if (buckets[i] == 0) break;
        }
        
        return false;
    }
    
    public boolean partitionKSubsets(int[] arr, int k) {
        int totalSum = 0;
        for (int num : arr) totalSum += num;
        
        if (totalSum % k != 0) return false;
        
        int target = totalSum / k;
        
        Integer[] boxedArr = Arrays.stream(arr).boxed().toArray(Integer[]::new);
        Arrays.sort(boxedArr, Collections.reverseOrder());
        arr = Arrays.stream(boxedArr).mapToInt(Integer::intValue).toArray();
        
        if (arr[0] > target) return false;
        
        int[] buckets = new int[k];
        return canPartition(arr, 0, buckets, k, target);
    }
}

Complexity Analysis

  • Time Complexity: O(k^n) - Each element can go to k buckets
  • Space Complexity: O(n) - Recursion stack

Approach 2: Optimal (Fill One Bucket at a Time)

Intuition

Instead of assigning elements to buckets, fill one bucket at a time. When a bucket is full (equals target), start filling the next bucket.

Algorithm

  1. Use bitmask to track used elements
  2. Fill buckets one at a time
  3. When current bucket reaches target, start new bucket
  4. When k buckets filled, return true
java
import java.util.*;

class Solution {
    public boolean backtrack(int[] arr, int bucketNum, int bucketSum, int startIndex,
                             int target, int k, boolean[] used) {
        if (bucketNum == k) {
            return true;
        }
        
        if (bucketSum == target) {
            return backtrack(arr, bucketNum + 1, 0, 0, target, k, used);
        }
        
        for (int i = startIndex; i < arr.length; i++) {
            if (used[i] || bucketSum + arr[i] > target) continue;
            
            if (i > 0 && arr[i] == arr[i - 1] && !used[i - 1]) continue;
            
            used[i] = true;
            
            if (backtrack(arr, bucketNum, bucketSum + arr[i], i + 1, target, k, used)) {
                return true;
            }
            
            used[i] = false;
            
            if (bucketSum == 0) break;
        }
        
        return false;
    }
    
    public boolean partitionKSubsets(int[] arr, int k) {
        int totalSum = 0;
        for (int num : arr) totalSum += num;
        
        if (totalSum % k != 0) return false;
        
        int target = totalSum / k;
        
        Integer[] boxedArr = Arrays.stream(arr).boxed().toArray(Integer[]::new);
        Arrays.sort(boxedArr, Collections.reverseOrder());
        arr = Arrays.stream(boxedArr).mapToInt(Integer::intValue).toArray();
        
        if (arr[0] > target) return false;
        
        boolean[] used = new boolean[arr.length];
        return backtrack(arr, 0, 0, 0, target, k, used);
    }
}

Complexity Analysis

  • Time Complexity: O(k × 2^n) - For each of k buckets, try 2^n subsets
  • Space Complexity: O(n) - Recursion and used array

Approach 3: Bitmask DP with Memoization

Complexity Analysis

  • Time Complexity: O(n × 2^n) - With memoization
  • Space Complexity: O(2^n) - Memoization table

Key Takeaways

  1. Early Checks: total % k != 0 or max element > target means impossible
  2. Sorting: Descending order for early pruning
  3. Bucket Approach: Two strategies - assign elements or fill buckets
  4. Empty Bucket Pruning: If empty bucket fails, other empty buckets will too
  5. LeetCode #698: Partition to K Equal Sum Subsets
  6. Related: Partition Equal Subset Sum (k=2)