698. Partition to K Equal Sum Subsets


Problem Description

In this LeetCode problem, we are tasked with determining whether a given array of integers nums can be divided into k non-empty subsets such that the sum of the numbers in each subset is equal. This problem is a variation of the classic partition problem which is known to be NP-complete, implying there is no known polynomial-time solution to solve it for the worst-case scenario. However, through clever approaches and optimizations, we can handle moderately sized inputs effectively.

To clarify:

  • A non-empty subset means each subset must contain at least one integer from nums.
  • All k subsets must have the same sum, and together they must include all elements from nums.

If such a division is possible, we should return true, otherwise false.

Intuition

The intuition behind solving this problem generally involves backtracking, a brute-force technique where we build subsets element by element and backtrack if we reach an invalid state. However, we can optimize our search with a few strategies:

  1. Pre-calculate subset sum (s): Since we aim for each of the k subsets to have an equal sum, we first compute the total sum of the array and divide it by k to find the target sum for each subset. If this division leaves a remainder (mod), we can immediately return false, because it means the total sum cannot be evenly divided into k parts.

  2. Sort and reverse: By sorting nums in descending order, we try larger numbers first. This is helpful because if larger numbers can't be placed into any subset, we'll fail faster and backtrack earlier, without wasting time on trying to fit smaller numbers first.

  3. Skip duplicate combinations: To avoid evaluating the same combination multiple times, we compare the current subset sum (cur[j]) with the previous one (cur[j - 1]). If they're equal and we're about to consider the same number again, we continue with the loop to try a different placement.

  4. State space tree: Use a depth-first search (DFS) backtracking algorithm to explore each state. The function dfs(i) takes an index i and tries to put nums[i] into each subset one by one. It recurses to check if adding nums[i] to a subset will lead to a solution (dfs(i + 1)).

  5. Backtracking: If adding nums[i] to the current subset doesn't lead to a solution or if the sum of the current subset exceeds the desired subset sum (s), we backtrack by removing nums[i] from cur[j] and try a different subset by continuing in the loop.

  6. Termination: The recursion terminates when all numbers are considered (i == len(nums)), at which point if we have not concluded to a contradiction, we return true indicating a successful partition.

Using these optimizations, we can significantly reduce the number of combinations we need to try and therefore solve the problem more efficiently than the naive brute-force approach.

Learn more about Memoization, Dynamic Programming, Backtracking and Bitmask patterns.

Solution Approach

The solution for the problem is implemented as a method canPartitionKSubsets(nums, k) in a Python class named Solution. Here's a step-by-step walkthrough of how the code works:

  1. Check divisibility: The sum of all elements in nums is divided by k to check if it's possible to partition the array into k subsets (s, mod = divmod(sum(nums), k)). If there's a remainder (mod is not 0), we return false immediately because an even partition is not possible.

  2. Data structure: A list (cur) is created with k zeros (cur = [0] * k), where each index represents the current sum of the corresponding subset.

  3. Sorting: The input list nums is sorted in reverse order (nums.sort(reverse=True)). This allows us to consider larger elements first, which is beneficial for early backtracking.

  4. Depth-first search (DFS): The core of the solution is a recursive function dfs(i) which attempts to find the correct placement for nums[i] in one of the subsets. The index i runs from 0 to len(nums) - 1.

  5. Termination condition: If i equals the length of nums (i == len(nums)), all items have been placed into some subset, and we return true, as we've found a valid partition.

  6. Subset iteration: Inside the dfs function, we loop over each subset (for j in range(k):), and try to place nums[i] into the jth subset.

  7. Duplicates check: To avoid placing the same combination again, we check if the current subset sum (cur[j]) is the same as the subset sum at the previous index (cur[j-1]). If it is, and j is not the first subset, we skip this iteration (continue) to prevent redundant work.

  8. Placement trial: We add nums[i] to the jth subset (cur[j] += nums[i]) and check if the new sum is less than or equal to the target sum (s).

  9. Recursion: If the new sum is valid, we call dfs(i + 1) to try and place the next element.

  10. Backtracking: If placing nums[i] into the jth subset doesn't lead to a solution, or the new sum exceeds the target, we backtrack by removing nums[i] from cur[j] (cur[j] -= nums[i]) and try a different subset in the next iteration.

  11. Final call and return statement: The last line of the method (return dfs(0)) starts the recursive process and returns the final answer, indicating whether a valid partition into k subsets with equal sums exists.

This algorithm employs backtracking, a form of DFS, to explore the solution space, and utilizes a few optimizations: sorting for fast backtracking, and duplicate combination skipping to reduce the state space.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Which of the following shows the order of node visit in a Breadth-first Search?

Example Walkthrough

Let's walk through an example to illustrate the solution approach with an array nums and an integer k. Suppose nums = [4, 3, 2, 3, 5, 2, 1] and k = 4. We need to find out if it's possible to partition this array into 4 subsets with equal sums.

  1. Check divisibility: First, we find the sum of all elements in nums, which is 4+3+2+3+5+2+1 = 20. Then we divide 20 by k (20 / 4) and find that each subset must sum up to 5 (s = 5), with no remainder, so an even partition is possible in theory.

  2. Data structure: We create a list cur = [0, 0, 0, 0] with 4 zeros, representing the sum of subsets we will form.

  3. Sorting: Sort nums in reverse to get [5, 4, 3, 3, 2, 2, 1]. This allows us to fill the subsets starting with the largest numbers.

  4. Depth-first search (DFS): We define our DFS to try placing each element in nums into one of the subsets.

  5. Termination condition: If we reach the end of nums, placing all elements without contradictions, we can return true.

  6. Subset iteration: We start with nums[0], which is 5, and try to place it in the first subset. cur[0] becomes 5, which is equal to our target subset sum (s), so the first subset is complete.

  7. Duplicates check: As we iterate through the rest of nums, we make sure not to create subsets with the same sum in the same manner twice.

  8. Placement trial: We take the next number, 4, and try adding it to our subsets, checking if the sum stays less than or equal to s. It won't fit in the first subset as it's already full, but it fits in the second (subset sum becomes 4).

  9. Recursion: We iterate with the next element, nums[2], and follow the previous steps. We keep trying until all elements are either placed or no placement is possible.

  10. Backtracking: If at any time, the current number doesn't fit in any subset, we backtrack, removing the last placed number and trying it in the next subset.

  11. Final call and return statement: At the end, we reach a valid partition where each subset sums to 5 with subsets being [5], [4, 1], [3, 2], [3, 2] or any permutation of these. The function will now return true, and our recursive process concludes that a valid partition is indeed possible.

In this example, with the help of sorting and careful placement of numbers, we minimized the number of operations required to find a valid partition. This demonstrates how the DFS backtracking algorithm efficiently solves the problem.

Solution Implementation

1from typing import List
2
3class Solution:
4    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
5        # Helper function that tries to assign elements to subsets recursively
6        def search(subset_sums, index):
7            # If we've traversed all the numbers, check if all subsets are equal to target sum
8            if index == len(nums):
9                return True
10            # Try adding the current number to each subset in turn and recurse
11            for j in range(k):
12                # Optimization: Skip if the previous subset is identical to the current
13                if j > 0 and subset_sums[j] == subset_sums[j - 1]:
14                    continue
15                subset_sums[j] += nums[index]
16                # Recurse only if the current subset sum is less than or equal to target subset sum
17                if subset_sums[j] <= target_sum and search(subset_sums, index + 1):
18                    return True
19                # Backtrack: remove the number from the current subset and try the next one
20                subset_sums[j] -= nums[index]
21            # If no valid assignment found, return False
22            return False
23
24        # Calculate the target subset sum and check if it is possible to divide
25        total_sum, remainder = divmod(sum(nums), k)
26        if remainder:
27            return False
28      
29        # Initialize k subset sums as 0
30        subset_sums = [0] * k
31        # Sort numbers in descending order to optimize the process
32        nums.sort(reverse=True)
33      
34        # Call the recursive search function starting with index 0
35        return search(subset_sums, 0)
36
1import java.util.Arrays;
2
3public class Solution {
4    private int[] numbers; // Array of input numbers.
5    private int[] subsetSums; // Array to hold the sum of the current subsets.
6    private int targetSubsetSum; // The desired sum of each subset.
7  
8    // Method to check if the numbers can be partitioned into k subsets with equal sum.
9    public boolean canPartitionKSubsets(int[] nums, int k) {
10        // Sum all the values in nums array.
11        int totalSum = Arrays.stream(nums).sum();
12      
13        // If totalSum is not divisible by k, we cannot partition it into k subsets of equal sum.
14        if (totalSum % k != 0) {
15            return false;
16        }
17      
18        // The target sum for each subset is the total sum divided by k.
19        this.targetSubsetSum = totalSum / k;
20      
21        // Sort the array in ascending order to optimize backtracking (largest elements are placed first).
22        Arrays.sort(nums);
23      
24        // Reverse the array to have the largest elements at the beginning.
25        for (int left = 0, right = nums.length - 1; left < right; left++, right--) {
26            int temp = nums[left];
27            nums[left] = nums[right];
28            nums[right] = temp;
29        }
30      
31        // Initialize member variables.
32        this.numbers = nums;
33        this.subsetSums = new int[k];
34      
35        // Start recursive DFS to find the solution.
36        return dfs(nums.length - 1);
37    }
38  
39    // Recursive method using Depth-First Search to backtrack and find whether a valid partition exists.
40    private boolean dfs(int currentIndex) {
41        // If we've gone through all the elements, check if all subset sums are equal to the target sum.
42        if (currentIndex < 0) {
43            for (int sum : subsetSums) {
44                if (sum != targetSubsetSum) {
45                    return false;
46                }
47            }
48            return true;
49        }
50      
51        // Try to place the current element into each subset one by one.
52        for (int j = 0; j < subsetSums.length; ++j) {
53            // Skip duplicate situations where previous and current subsetSums are the same value (to avoid redundant work).
54            if (j > 0 && subsetSums[j] == subsetSums[j - 1]) {
55                continue;
56            }
57          
58            // Add current element to the subset sum.
59            subsetSums[j] += numbers[currentIndex];
60          
61            // If the new subset sum is less than or equal to the target,
62            // continue DFS with the next smaller element.
63            if (subsetSums[j] <= targetSubsetSum && dfs(currentIndex - 1)) {
64                return true;
65            }
66          
67            // Backtrack: remove current element from the subset sum.
68            subsetSums[j] -= numbers[currentIndex];
69          
70            // If the sum is zero after backtracking, no need to try placing the current element into later buckets.
71            if (subsetSums[j] == 0) break;
72        }
73      
74        // If no valid placement was found, return false.
75        return false;
76    }
77}
78
1#include <vector>
2#include <numeric>
3#include <algorithm>
4#include <functional>
5
6class Solution {
7public:
8    bool canPartitionKSubsets(std::vector<int>& nums, int k) {
9        // Calculate the sum of all numbers
10        int totalSum = std::accumulate(nums.begin(), nums.end(), 0);
11      
12        // If totalSum is not divisible by k, we can't partition it into k subsets.
13        if (totalSum % k) {
14            return false;
15        }
16      
17        // The target sum for each subset.
18        int targetSum = totalSum / k;
19      
20        // Sort numbers in descending order to help with the backtracking process.
21        std::sort(nums.begin(), nums.end(), std::greater<int>());
22
23        // Vector to store the current sum of each subset.
24        std::vector<int> subsetSums(k, 0);
25      
26        // Define a depth-first search lambda to explore the possible partitions.
27        std::function<bool(int)> dfs = [&](int index) {
28            // If we have reached the end of nums, check if we successfully partitioned.
29            if (index == nums.size()) {
30                // All k subsets have our target sum.
31                return true;
32            }
33          
34            // Iterate through each subset
35            for (int j = 0; j < k; ++j) {
36                // To avoid redundant work, if the previous subset has the same sum
37                // as the current one, we don't need to try nums[index] in this subset.
38                if (j > 0 && subsetSums[j] == subsetSums[j - 1]) {
39                    continue;
40                }
41                // Add the current number to the current subset sum.
42                subsetSums[j] += nums[index];
43              
44                // If the new sum does not exceed the target, and dfs succeeds, return true.
45                if (subsetSums[j] <= targetSum && dfs(index + 1)) {
46                    return true;
47                }
48              
49                // Backtrack: remove the current number from the current subset sum.
50                subsetSums[j] -= nums[index];
51            }
52            // If no valid partitioning is found, return false.
53            return false;
54        };
55
56        // Start the depth-first search with the first number.
57        return dfs(0);
58    }
59};
60
1function canPartitionKSubsets(nums: number[], k: number): boolean {
2    // Calculates the sum of all elements in nums
3    let totalSum = nums.reduce((accumulator, currentValue) => accumulator + currentValue);
4
5    // If the total sum is not divisible by k, it's not possible to partition the array
6    if (totalSum % k !== 0) {
7        return false;
8    }
9
10    // The target sum for each subset is the total sum divided by k
11    let targetSubsetSum = totalSum / k;
12
13    // Sort the array in ascending order to optimize the algorithm
14    nums.sort((a, b) => a - b);
15
16    // Get the length of nums
17    const n = nums.length;
18
19    // Initialize a boolean array to keep track of possible subset sums
20    const canFormSubsetSum = new Array(1 << n).fill(false);
21    canFormSubsetSum[0] = true;
22
23    // Initialize an array to store the current subset sum
24    const currentSubsetSum = new Array(n).fill(0);
25
26    // Iterate over all possible subsets
27    for (let i = 0; i < 1 << n; ++i) {
28        // If current subset cannot form the sum, skip to the next iteration
29        if (!canFormSubsetSum[i]) {
30            continue;
31        }
32
33        // Try to add each number to the current subset to form a new subset
34        for (let j = 0; j < n; ++j) {
35            // Skip if adding the number exceeds the target subset sum
36            if (currentSubsetSum[i] + nums[j] > targetSubsetSum) {
37                break;
38            }
39
40            // Check if the number is not already used in the current subset
41            if (((i >> j) & 1) === 0) {
42                const newSubsetIndex = i | (1 << j);
43                canFormSubsetSum[newSubsetIndex] = true; // Mark the new subset as possible
44                currentSubsetSum[newSubsetIndex] = (currentSubsetSum[i] + nums[j]) % targetSubsetSum;
45            }
46        }
47    }
48
49    // Check if there is a combination that includes all numbers forming k subsets
50    return canFormSubsetSum[(1 << n) - 1];
51}
52

Time and Space Complexity

The given Python code is trying to determine if it's possible to partition the input list nums into k subsets each with the same sum. The approach is using backtracking to explore all possible ways to create these partitions.

Time Complexity

The time complexity of this function can be quite high due to the nature of the problem, which is NP-complete.

The dfs function tries to place each number into one of the k subsets in a depth-first search manner. Given that there are n numbers, for the first number there are k choices, for the second number there are k choices, and so on, leading to a total of k^n possibilities in the worst case. However, there are some optimizations:

  1. The for loop in the dfs function has pruning: it skips placing a number in a subset if the current subset has the same sum as the previous subset, which could reduce the number of explored branches.
  2. Sorting the array in reverse order helps in an earlier detection of sums that exceed the target subset sum s, potentially reducing the number of recursive calls.

With such optimizations, the upper bound on the number of calls to dfs is less than k^n, but it is hard to quantify the exact impact without a specific case analysis.

Therefore, the time complexity is better represented as O(k^n), keeping in mind that this is a simplified upper bound and actual performance may be significantly better due to the mentioned optimizations.

Space Complexity

The space complexity is determined by:

  1. The depth of the recursion, which in the worst case is O(n) since dfs might go as deep as the number of elements in the list.
  2. The space used by the list cur which is O(k).

Since these two factors are not dependent on each other, we add them up to get the total space complexity: O(n + k).

Remember that the recursion call stack space (for dfs) has a space complexity of O(n) because the depth of the recursion tree can go up to the number of elements in nums. The additional O(k) space is for maintaining the current sums of the k subsets.

Learn more about how to find time and space complexity quickly using problem constraints.


Fast Track Your Learning with Our Quick Skills Quiz:

Which of the following is a min heap?


Recommended Readings


Got a question? Ask the Monster Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.

Tired of the LeetCode Grind?

Our structured approach teaches you the patterns behind problems, so you can confidently solve any challenge. Get started now to land your dream tech job.

Get Started

🪄