Partition to K Equal Sum Subsets
Problem
Can the array be partitioned into k non-empty subsets with equal sum?
nums=[4,3,2,3,5,2,1] k=4Truedef can_partition_k_subsets(nums, k):
s = sum(nums)
if s % k != 0: return False
target = s // k; nums.sort(reverse=True)
if nums[0] > target: return False
bucket = [0] * k
def go(i):
if i == len(nums): return all(b == target for b in bucket)
seen = set()
for j in range(k):
if bucket[j] in seen: continue
if bucket[j] + nums[i] > target: continue
seen.add(bucket[j]); bucket[j] += nums[i]
if go(i + 1): return True
bucket[j] -= nums[i]
if bucket[j] == 0: break
return False
return go(0)function canPartitionKSubsets(nums, k) {
const s = nums.reduce((a, b) => a + b, 0);
if (s % k !== 0) return false;
const target = s / k; nums.sort((a, b) => b - a);
if (nums[0] > target) return false;
const bucket = new Array(k).fill(0);
const go = i => {
if (i === nums.length) return bucket.every(b => b === target);
const seen = new Set();
for (let j = 0; j < k; j++) {
if (seen.has(bucket[j])) continue;
if (bucket[j] + nums[i] > target) continue;
seen.add(bucket[j]); bucket[j] += nums[i];
if (go(i + 1)) return true;
bucket[j] -= nums[i];
if (bucket[j] === 0) break;
}
return false;
};
return go(0);
}boolean canPartitionKSubsets(int[] nums, int k) {
int s = 0; for (int x : nums) s += x;
if (s % k != 0) return false;
int target = s / k;
Integer[] a = Arrays.stream(nums).boxed().toArray(Integer[]::new);
Arrays.sort(a, Comparator.reverseOrder());
if (a[0] > target) return false;
int[] bucket = new int[k];
return go(a, 0, bucket, target);
}
boolean go(Integer[] a, int i, int[] bucket, int target) {
if (i == a.length) { for (int b : bucket) if (b != target) return false; return true; }
Set seen = new HashSet<>();
for (int j = 0; j < bucket.length; j++) {
if (seen.contains(bucket[j])) continue;
if (bucket[j] + a[i] > target) continue;
seen.add(bucket[j]); bucket[j] += a[i];
if (go(a, i + 1, bucket, target)) return true;
bucket[j] -= a[i];
if (bucket[j] == 0) break;
}
return false;
} bool go(vector& a, int i, vector& bucket, int target) {
if (i == (int)a.size()) { for (int b : bucket) if (b != target) return false; return true; }
set seen;
for (int j = 0; j < (int)bucket.size(); j++) {
if (seen.count(bucket[j])) continue;
if (bucket[j] + a[i] > target) continue;
seen.insert(bucket[j]); bucket[j] += a[i];
if (go(a, i + 1, bucket, target)) return true;
bucket[j] -= a[i];
if (bucket[j] == 0) break;
}
return false;
}
bool canPartitionKSubsets(vector& nums, int k) {
int s = 0; for (int x : nums) s += x;
if (s % k != 0) return false;
int target = s / k;
sort(nums.begin(), nums.end(), greater());
if (nums[0] > target) return false;
vector bucket(k, 0);
return go(nums, 0, bucket, target);
} Explanation
We want to split the numbers into k groups that all add up to the same amount. First we figure out what that amount must be: add everything up and divide by k. If it does not divide evenly, the answer is immediately False.
Each group is a bucket, and we have k empty buckets with a goal of target = sum / k. We then try to place numbers into buckets one by one using recursion with backtracking: drop a number into a bucket, recurse to place the next number, and if things go wrong, take it back out and try a different bucket.
Two clever tricks make this fast. We sort the numbers largest first so big values are placed early (they fail faster), and we keep a seen set so we never try two buckets that currently hold the same amount — putting a number in either would look identical.
The line if bucket[j] == 0: break is another shortcut: if placing a number in an empty bucket did not lead to a solution, no other empty bucket will either, so we stop early.
Example: nums=[4,3,2,3,5,2,1], k=4. The total is 20, so target = 5. The numbers can be grouped as 5, 4+1, 3+2, 3+2 — four buckets all summing to 5 — so the answer is True.