Partition in k equal subsets

The idea of the algorithm is not to test each possible combination for every bucket but rather to try each possible placement for every element.

1) We start with a list of the bucket sums initialized to 0

2) Then we try to place the first element (by decreasing order)

a) iterate over every bucket

b) place the current element in the current bucket

c) add its value to the sum of the bucket

d) check if the updated sum of the bucket is smaller or equal to the target

e.i) if yes apply the same function to the next element with the updated sums of the buckets

e.ii) if no just remove the element from the bucket and try to place it in the next bucket

f) if every element have been placed, just check if every bucket have the same sum equal to the target

g.i) if yes we are done

g.ii) otherwise this combination of placements is not good

h) remove the current element from the bucket and try to place it in another bucket

i) if no placement works for the current element hence the current placements of previous elements is not ok and we need to backtrack to place the precedent element in another bucket

def canPartitionKSubsets(nums: List[int], k: int) -> bool:
    buckets = [0]*k # initialize each bucket sum to 0
    target = sum(nums) // k

    # We want to try placing larger numbers first
    nums.sort(reverse=True)

    # dfs determines which bucket to put the 'current element' (nums[idx]) into
    # it is dfs as we try all possible combination for one move before testing another move
    def dfs(idx):
        # If we've placed all of the items, we're done
        if idx == len(nums):
            return set(buckets) == set([target]) # if each bucket value is target return True otherwise False

        # For the current element iterate over each bucket
        for i in range(k):
            # Add the current element to it
            buckets[i] += nums[idx]

            # If it's a valid placement we now try to place the remaining elements starting with the next one.
            # If we can obtain a correct partition with this move return True
            if buckets[i] <= target and dfs(idx + 1):
                return True

            # Otherwise remove the current element from the ith bucket and try the next one. 
            buckets[i] -= nums[idx]

            # if this bucket was originally empty and we could not place this element in it we do not need
            # to test next buckets as they are also empty (it's an iterative process over the buckets)
            if buckets[i] == 0:
                break

        # If we couldn't place the current element anywhere that leads to a valid solution
        # we will need to backtrack and try something else with precedent elements. 
        return False

    # Start by trying to place first element ie nums[0]
    return dfs(0)