698. Partition To K Equal Sum Subsets
class Solution:
"""🎲 Partition to K Equal Sum Subsets.
Problem:
--------
Given an integer array `nums` (positive values, possibly repeated) and an integer `k`,
return whether `nums` can be divided into `k` non-empty subsets with equal sums. Every
element must be used exactly once.
Approach:
---------
Each subset must sum to `target = sum(nums) // k`, so a necessary (not sufficient)
condition is `sum(nums) % k == 0`; likewise no element may exceed `target`.
Backtracking over the elements: keep `k` bucket sums and, for each element in turn,
try placing it into each bucket.
- Complete: every element is placed. No bucket exceeds `target` and the buckets sum to
`k * target`, so each holds exactly `target` (and is non-empty, since `target > 0`).
- Make move: add `nums[idx]` to a bucket, then recurse on `idx + 1`.
- Unmake move: subtract it again. Return as soon as one branch succeeds.
- Prune:
- Skip a bucket that would overflow `target`.
- Skip a bucket whose current sum was already tried for this element: buckets with
equal sums are interchangeable, so placing the element there gives a symmetric state
that fails the same way. In particular all empty buckets are equivalent, so the first
element is tried in only one of them.
- Sort `nums` in decreasing order: large elements have the fewest valid buckets, so
dead ends surface near the root instead of deep in the tree.
Complexity:
-----------
- Time: O(k^n) worst case, each of the n elements tried in up to k buckets; pruning cuts
this sharply in practice. The O(n log n) sort is dominated.
- Space: O(n * k) auxiliary, recursion depth n with a `seen` set of up to k sums per frame,
plus O(n) for the sorted copy and O(k) for the buckets.
"""
def canPartitionKSubsets(self, nums: list[int], k: int) -> bool:
total_sum = sum(nums)
if total_sum % k != 0:
return False
target = total_sum // k
# Largest first, so dead ends are found near the root
nums = sorted(nums, reverse=True)
if nums[0] > target:
return False
buckets = [0] * k
def place(idx: int) -> bool:
"""Try to place nums[idx] into one of the buckets."""
if idx == len(nums):
return True
seen: set[int] = set()
for bucket_idx in range(k):
# Skip overflowing buckets and buckets equivalent to one already tried
if buckets[bucket_idx] + nums[idx] > target or buckets[bucket_idx] in seen:
continue
seen.add(buckets[bucket_idx])
buckets[bucket_idx] += nums[idx]
if place(idx + 1):
return True
buckets[bucket_idx] -= nums[idx]
return False
return place(0)