698. 划分为k个相等的子集
题目描述
输入一个正整数数组 nums 和一个正整数 k。
判断能否把 nums 中的全部元素分到 k 个非空的组里,每个元素恰好属于一组,并且每组元素之和都相同。能则返回 true,否则返回 false。
输入:nums = [4,3,2,3,5,2,1], k = 4 输出:true
解释:可以分成 (5)、(1,4)、(2,3)、(2,3),每组和都是 5。
输入:nums = [1,2,3,4], k = 3 输出:false
解释:总和 10 不能被 3 整除。
数据范围
1 <= k <= nums.length <= 160 < nums[i] < 10000- 每个值出现的次数在
[1, 4]之间。
解题思路:回溯(按桶放置 + 剪枝)
每个子集的和必须是 side = sum / k,不能整除就直接失败。把 k 个子集看成 k 个桶,按从大到小的顺序把每个数放进某个桶,放入后桶内和不能超过 side。大数先放,失败会更早暴露。若当前桶与前一个桶的和相同,放进去的结果是对称的,前一个桶已经试过,直接跳过。所有数放完时每个桶恰好等于 side。
为什么正确
进入 dfs(i) 时,前 i 个数已经分配,每个桶的和都不超过 side。
复杂度分析
时间最坏 O(k^n),排序和对称剪枝后实际很快;空间 O(k + n)。
易错点
回溯后忘记把数从桶里拿出来;没有排序和对称剪枝时会超时。还要注意总和不能被 k 整除或最大数超过 side 时直接返回 False。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。
class Solution:
def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
total = sum(nums)
if total % k:
return False
side = total // k
nums.sort(reverse=True)
if nums[0] > side:
return False
buckets = [0] * k
def dfs(i):
if i == len(nums):
return True
for b in range(k):
if buckets[b] + nums[i] > side:
continue
if b > 0 and buckets[b] == buckets[b - 1]:
continue
buckets[b] += nums[i]
if dfs(i + 1):
return True
buckets[b] -= nums[i]
return False
return dfs(0)class Solution {
int[] nums, buckets;
int k, side;
public boolean canPartitionKSubsets(int[] nums, int k) {
int total = 0;
for (int v : nums) total += v;
if (total % k != 0) {
return false;
}
side = total / k;
Arrays.sort(nums);
for (int l = 0, r = nums.length - 1; l < r; l++, r--) { int t = nums[l]; nums[l] = nums[r]; nums[r] = t; }
if (nums[0] > side) {
return false;
}
buckets = new int[k];
this.nums = nums;
this.k = k;
return dfs(0);
}
boolean dfs(int i) {
if (i == nums.length) {
return true;
}
for (int b = 0; b < k; b++) {
if (buckets[b] + nums[i] > side) {
continue;
}
if (b > 0 && buckets[b] == buckets[b - 1]) {
continue;
}
buckets[b] += nums[i];
if (dfs(i + 1)) {
return true;
}
buckets[b] -= nums[i];
}
return false;
}
}class Solution {
public:
bool canPartitionKSubsets(vector<int>& nums, int k) {
int total = accumulate(nums.begin(), nums.end(), 0);
if (total % k) {
return false;
}
int side = total / k;
sort(nums.rbegin(), nums.rend());
if (nums[0] > side) {
return false;
}
vector<int> buckets(k, 0);
function<bool(int)> dfs = [&](int i) -> bool {
if (i == (int)nums.size()) {
return true;
}
for (int b = 0; b < k; b++) {
if (buckets[b] + nums[i] > side) {
continue;
}
if (b > 0 && buckets[b] == buckets[b - 1]) {
continue;
}
buckets[b] += nums[i];
if (dfs(i + 1)) {
return true;
}
buckets[b] -= nums[i];
}
return false;
};
return dfs(0);
}
};func canPartitionKSubsets(nums []int, k int) bool {
total := 0
for _, v := range nums {
total += v
}
if total%k != 0 {
return false
}
side := total / k
sort.Sort(sort.Reverse(sort.IntSlice(nums)))
if nums[0] > side {
return false
}
buckets := make([]int, k)
var dfs func(i int) bool
dfs = func(i int) bool {
if i == len(nums) {
return true
}
for b := 0; b < k; b++ {
if buckets[b]+nums[i] > side {
continue
}
if b > 0 && buckets[b] == buckets[b-1] {
continue
}
buckets[b] += nums[i]
if dfs(i + 1) {
return true
}
buckets[b] -= nums[i]
}
return false
}
return dfs(0)
}var canPartitionKSubsets = function(nums, k) {
const total = nums.reduce((a, b) => a + b, 0);
if (total % k !== 0) {
return false;
}
const side = total / k;
nums.sort((a, b) => b - a);
if (nums[0] > side) {
return false;
}
const buckets = new Array(k).fill(0);
const dfs = (i) => {
if (i === nums.length) {
return true;
}
for (let b = 0; b < k; b++) {
if (buckets[b] + nums[i] > side) {
continue;
}
if (b > 0 && buckets[b] === buckets[b - 1]) {
continue;
}
buckets[b] += nums[i];
if (dfs(i + 1)) {
return true;
}
buckets[b] -= nums[i];
}
return false;
};
return dfs(0);
};function canPartitionKSubsets(nums: number[], k: number): boolean {
const total = nums.reduce((a, b) => a + b, 0);
if (total % k !== 0) {
return false;
}
const side = total / k;
nums.sort((a, b) => b - a);
if (nums[0] > side) {
return false;
}
const buckets: number[] = new Array(k).fill(0);
const dfs = (i: number): boolean => {
if (i === nums.length) {
return true;
}
for (let b = 0; b < k; b++) {
if (buckets[b] + nums[i] > side) {
continue;
}
if (b > 0 && buckets[b] === buckets[b - 1]) {
continue;
}
buckets[b] += nums[i];
if (dfs(i + 1)) {
return true;
}
buckets[b] -= nums[i];
}
return false;
};
return dfs(0);
}其他解法
回溯以桶为单位尝试放置,靠剪枝减少搜索;状态压缩把“哪些数已用”当作状态,去掉重复搜索,最坏情况可控。
解法 2:状态压缩动态规划 时间 O(n·2ⁿ),空间 O(2ⁿ)
用二进制掩码 mask 表示已经用掉的数。按顺序填桶,一个桶填满再填下一个,所以当前桶里的和只由 mask 决定,等于已用数之和模 side。dp[mask] 表示这个状态能否合法到达。从可达状态出发,尝试加入一个未用的数,只要不超过 side 就转移。数组先排序,放不下时可以直接结束枚举。全集可达即为真。
每个状态只算一次,最坏复杂度有保证,不依赖剪枝效果。代价是 2ⁿ 的表,只适合 n ≤ 16 这样的小规模。
class Solution:
def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
total = sum(nums)
if total % k:
return False
side = total // k
nums.sort()
if nums[-1] > side:
return False
n = len(nums)
full = (1 << n) - 1
dp = [False] * (1 << n)
cur = [0] * (1 << n)
dp[0] = True
for mask in range(1 << n):
if not dp[mask]:
continue
for i in range(n):
if cur[mask] + nums[i] > side:
break
if mask >> i & 1:
continue
nxt = mask | 1 << i
if not dp[nxt]:
dp[nxt] = True
cur[nxt] = (cur[mask] + nums[i]) % side
return dp[full]class Solution {
public boolean canPartitionKSubsets(int[] nums, int k) {
int total = 0;
for (int v : nums) total += v;
if (total % k != 0) return false;
int side = total / k, n = nums.length;
Arrays.sort(nums);
if (nums[n - 1] > side) return false;
boolean[] dp = new boolean[1 << n];
int[] cur = new int[1 << n];
dp[0] = true;
for (int mask = 0; mask < (1 << n); mask++) {
if (!dp[mask]) continue;
for (int i = 0; i < n; i++) {
if (cur[mask] + nums[i] > side) break;
if ((mask >> i & 1) != 0) continue;
int nxt = mask | 1 << i;
if (!dp[nxt]) {
dp[nxt] = true;
cur[nxt] = (cur[mask] + nums[i]) % side;
}
}
}
return dp[(1 << n) - 1];
}
}class Solution {
public:
bool canPartitionKSubsets(vector<int>& nums, int k) {
int total = accumulate(nums.begin(), nums.end(), 0);
if (total % k) return false;
int side = total / k, n = nums.size();
sort(nums.begin(), nums.end());
if (nums.back() > side) return false;
vector<bool> dp(1 << n, false);
vector<int> cur(1 << n, 0);
dp[0] = true;
for (int mask = 0; mask < (1 << n); mask++) {
if (!dp[mask]) continue;
for (int i = 0; i < n; i++) {
if (cur[mask] + nums[i] > side) break;
if (mask >> i & 1) continue;
int nxt = mask | 1 << i;
if (!dp[nxt]) {
dp[nxt] = true;
cur[nxt] = (cur[mask] + nums[i]) % side;
}
}
}
return dp[(1 << n) - 1];
}
};func canPartitionKSubsets(nums []int, k int) bool {
total := 0
for _, v := range nums {
total += v
}
if total%k != 0 {
return false
}
side, n := total/k, len(nums)
sort.Ints(nums)
if nums[n-1] > side {
return false
}
dp := make([]bool, 1<<n)
cur := make([]int, 1<<n)
dp[0] = true
for mask := 0; mask < 1<<n; mask++ {
if !dp[mask] {
continue
}
for i := 0; i < n; i++ {
if cur[mask]+nums[i] > side {
break
}
if mask>>i&1 == 1 {
continue
}
nxt := mask | 1<<i
if !dp[nxt] {
dp[nxt] = true
cur[nxt] = (cur[mask] + nums[i]) % side
}
}
}
return dp[1<<n-1]
}var canPartitionKSubsets = function(nums, k) {
const total = nums.reduce((a, b) => a + b, 0);
if (total % k !== 0) return false;
const side = total / k, n = nums.length;
nums.sort((a, b) => a - b);
if (nums[n - 1] > side) return false;
const dp = new Uint8Array(1 << n);
const cur = new Int32Array(1 << n);
dp[0] = 1;
for (let mask = 0; mask < (1 << n); mask++) {
if (!dp[mask]) continue;
for (let i = 0; i < n; i++) {
if (cur[mask] + nums[i] > side) break;
if ((mask >> i) & 1) continue;
const nxt = mask | (1 << i);
if (!dp[nxt]) {
dp[nxt] = 1;
cur[nxt] = (cur[mask] + nums[i]) % side;
}
}
}
return dp[(1 << n) - 1] === 1;
};function canPartitionKSubsets(nums: number[], k: number): boolean {
const total = nums.reduce((a, b) => a + b, 0);
if (total % k !== 0) return false;
const side = total / k, n = nums.length;
nums.sort((a, b) => a - b);
if (nums[n - 1] > side) return false;
const dp = new Uint8Array(1 << n);
const cur = new Int32Array(1 << n);
dp[0] = 1;
for (let mask = 0; mask < (1 << n); mask++) {
if (!dp[mask]) continue;
for (let i = 0; i < n; i++) {
if (cur[mask] + nums[i] > side) break;
if ((mask >> i) & 1) continue;
const nxt = mask | (1 << i);
if (!dp[nxt]) {
dp[nxt] = 1;
cur[nxt] = (cur[mask] + nums[i]) % side;
}
}
}
return dp[(1 << n) - 1] === 1;
}