Split numbers into k groups with equal sums: find the target, place the biggest numbers first, and skip buckets that repeat a sum.
▼The problem
LeetCode 698 (Medium). Given an integer array nums and an integer k, return true if nums can be divided into k non-empty subsets whose sums are all equal.
Examples (LeetCode's): nums = [4, 3, 2, 3, 5, 2, 1], k = 4 → true (sum 20, target 5: [5], [1, 4], [2, 3], [2, 3]; walked through in scene 6) and nums = [1, 2, 3, 4], k = 3 → false (sum 10 is not divisible by 3).
The solution
def canPartitionKSubsets(nums, k):
total = sum(nums)
if total % k: return False
target = total // k
nums.sort(reverse=True)
if nums[0] > target: return False
bins = [0] * k
def place(i):
if i == len(nums): return True
seen = set()
for b in range(k):
if bins[b] + nums[i] > target or bins[b] in seen:
continue
seen.add(bins[b])
bins[b] += nums[i]
if place(i + 1): return True
bins[b] -= nums[i]
if bins[b] == 0: break # failed in an empty bin
return False
return place(0)Transcript
Partition to K Equal Sum Subsets. Given a list of numbers and k, can you split them into k groups with equal sums?
Take four, three, two, three, five, two, one, with k equal to four. Twenty total, so each group makes five: five, one and four, two and three, two and three. True. But one, two, three, four into three groups totals ten, which three can't divide. False.
The naive way tries every number in every group: k to the n assignments. Seven numbers in four groups is over sixteen thousand.
Better: the target is the total over k. If it's not whole, or the biggest number is larger, stop. Sort big to small, and place each number where its group stays within the target. Two prunes help. If two groups hold the same sum, skip the second. If a number fails in an empty group, every empty group fails, so back up.
In code, each call places the next number, recurses, and takes it back on failure. A bitmask dynamic program, n times two to the n, also works.
Let's pack it. Five fills group one. Four starts group two. Three fits neither; empty groups are alike, so it opens just group three. The other three opens group four. The twos top up groups three and four, and one finishes group two. True.
The worst case is still exponential, but pruning cuts most branches. Space is order n plus k.
Find the target, sort big first, skip repeated sums. That's Partition to K Equal Sum Subsets.