Taro Logo

Partition to K Equal Sum Subsets

Medium
Google logo
Google
3 views
Topics:
ArraysDynamic ProgrammingBit Manipulation

Given an integer array nums and an integer k, determine if it's possible to divide this array into k non-empty subsets such that the sum of elements in each subset is equal.

Example 1:

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

Output: true

Explanation: The array can be divided into the following 4 subsets: (5), (1, 4), (2, 3), (2, 3), each having a sum of 5.

Example 2:

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

Output: false

Constraints:

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

Could you implement a function that solves this problem efficiently?

Solution


Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

  1. What are the constraints on the range of values in the input array, and what is the maximum size of the array? Can I assume the input array will only contain integers?
  2. Is it possible for k to be zero, one, or greater than the number of elements in the input array? What should the function return in those cases?
  3. If it's not possible to partition the array into k equal sum subsets, what should the function return (e.g., false, null, or an exception)?
  4. Are duplicate numbers allowed in the input array, and if so, does their presence affect the solution or time complexity significantly?
  5. If there are multiple valid partitions into k equal sum subsets, does the function need to return all of them, or is returning any one valid partition sufficient?

Brute Force Solution

Approach

The brute force approach to this problem involves trying every single possible way to divide the numbers into the required number of groups. We explore all combinations to see if any arrangement results in each group having the same sum. Think of it like shuffling cards and dealing them into piles to see if the piles have the same total value.

Here's how the algorithm would work step-by-step:

  1. First, we try putting the first number into the first group.
  2. Then, we try putting the first number into the second group, and so on, up to the last group.
  3. For each of these choices, we consider the second number and repeat the process: try putting it into each of the groups.
  4. We continue this process for every number, trying all possible group assignments.
  5. After assigning all numbers to groups, we check if all groups have the same sum.
  6. If the sums are all equal, we found a valid arrangement!
  7. If not, we go back and try a different arrangement until we have exhausted all possible combinations.
  8. If we have tried every possible combination and none of them resulted in groups with equal sums, then it's not possible to partition the numbers in that way.

Code Implementation

def can_partition_to_k_equal_sum_subsets_brute_force(
    numbers, number_of_subsets
):
    total_sum = sum(numbers)
    if total_sum % number_of_subsets != 0:
        return False

    target_sum = total_sum // number_of_subsets
    subsets = [0] * number_of_subsets

    def can_partition_recursive(index):
        if index == len(numbers):
            # Check if all subsets have the same sum
            return all(subset_sum == target_sum for subset_sum in subsets)

        for i in range(number_of_subsets):
            # Try adding the current number to each subset
            subsets[i] += numbers[index]

            if subsets[i] <= target_sum:
                if can_partition_recursive(index + 1):
                    return True

            # Backtrack: Undo the addition to explore other possibilities
            subsets[i] -= numbers[index]

        return False

    # Initiate the recursive process
    return can_partition_recursive(0)

Big(O) Analysis

Time Complexity
O(K^N)The brute force approach involves assigning each of the N numbers to one of the K subsets. For each of the N numbers, there are K choices of which subset to put it in. Therefore, there are K * K * ... * K (N times) which is K^N possible assignments of numbers to subsets. After each assignment, we need to check if the sums of all subsets are equal which takes O(K) time where K is the number of subsets. Thus, the total runtime complexity becomes O(K^N * K) which simplifies to O(K^N) since K is a constant factor multiplied with K^N.
Space Complexity
O(K^N)The brute force approach explores all possible group assignments for each number. Since each of the N numbers can be assigned to one of the K groups, the recursion depth can go up to N. At each level of the recursion, we are essentially storing the current assignment of numbers to groups in the call stack. In the worst-case scenario, we explore all K^N possible combinations of assignments which contributes to the exponential space complexity. Therefore, the space required for the recursion stack is proportional to K raised to the power of N.

Optimal Solution

Approach

The core idea is to see if it's even possible to split the numbers into equal sum groups before doing any real work. Then, we try to build these groups one at a time, stopping early if we realize we're on the wrong path.

Here's how the algorithm would work step-by-step:

  1. First, check if the numbers can actually be divided equally. If the total sum of the numbers is not divisible by the number of groups you want, it's impossible, and you're done.
  2. Calculate the target sum for each group. This is the total sum divided by the desired number of groups.
  3. Start trying to fill each group one at a time.
  4. For each group, consider the numbers and see if you can add some of them up to reach the target sum.
  5. If you find a combination of numbers that adds up to the target sum for a group, mark those numbers as 'used'.
  6. Repeat the process for the next group, making sure not to use the numbers that were already used in previous groups.
  7. If, at any point, you can't find a combination of unused numbers to reach the target sum for a group, backtrack. This means undo your previous choices and try a different path to build the groups.
  8. If you successfully fill all the groups with their respective numbers, then you've found a valid split.

Code Implementation

def canPartitionKSubsets(numbers, k):
    total_sum = sum(numbers)
    if total_sum % k != 0:
        return False

    target_sum_for_subset = total_sum // k
    number_of_elements = len(numbers)
    used_elements = [False] * number_of_elements

    def can_form_subset(current_subset_sum, elements_used,
                        start_index, subsets_formed):

        if subsets_formed == k:
            return True

        if current_subset_sum == target_sum_for_subset:
            # Move onto forming the next subset
            return can_form_subset(0, elements_used, 0, subsets_formed + 1)

        if current_subset_sum > target_sum_for_subset:
            return False

        for i in range(start_index, number_of_elements):
            if not elements_used[i]:
                elements_used[i] = True
                # Explore using this number in the current subset
                if can_form_subset(current_subset_sum + numbers[i], 
                                    elements_used, i + 1, subsets_formed):
                    return True
                
                # Backtrack if including the number didn't work
                elements_used[i] = False
        
        return False
    
    # If total_sum % k is 0, target sum is calculated
    return can_form_subset(0, used_elements, 0, 0)

Big(O) Analysis

Time Complexity
O(k * 2^n)The dominant factor is the recursive exploration of subsets. For each of the k groups, in the worst case, we may need to consider every possible subset of the n numbers to find a combination that sums to the target. Checking all subsets leads to 2^n possibilities. Since we potentially repeat this for each of the k groups, the overall time complexity becomes O(k * 2^n). The backtracking adds significantly to the calculations as well.
Space Complexity
O(N)The provided algorithm uses a 'used' array to keep track of which numbers have been assigned to a subset. This array has a size equal to the number of input numbers, denoted as N. Additionally, the recursive calls, in the worst-case scenario, can add up to a depth of N (if each number is considered individually before backtracking), leading to a call stack of size N. Thus, the auxiliary space used is proportional to N, resulting in a space complexity of O(N).

Edge Cases

CaseHow to Handle
Null or empty input arrayReturn False, as partitioning an empty set is impossible.
k is zero or negativeReturn False, as we cannot partition into a non-positive number of subsets.
k is greater than the number of elements in the arrayReturn False, as each subset must have at least one element.
Sum of array elements is not divisible by kReturn False, as equal sum subsets are impossible if the total sum is not evenly divisible.
Array contains negative numbersThe problem statement does not specify that all numbers must be non-negative, so the backtracking logic needs to account for negative values, and should return False if the subset target becomes negative.
Array contains large numbers, leading to integer overflow when calculating the sumUse a data type with a larger range (e.g., long) to store the sum of array elements.
All elements are zeroIf the total sum is zero and k is less than or equal to n, then a valid partition is possible.
A single element is greater than the target subset sumReturn False, because no valid partition can exist if one element is already greater than the target subset sum.