一、题目
给定一个整数数组nums
和一个正整数
k
,找出是否有可能把这个数组分成
k
个非空子集,其总和都相等
。
二、解题思路
在前面,我们已经讨论过一道与本题极为相似的问题,那一题目是要求使用所有火柴能否拼成一个正方形。与之相比,本题的差异在于我们需要判断是否能用所有火柴(即数组中的元素)拼成一个正K边形。之前可以视为对一棵四叉树进行深度优先搜索(DFS)遍历,而本题则可以类比为对一棵K叉树进行DFS遍历。两者的代码结构极为相近,因此我们可以借鉴前面的代码并进行适当修改来解决本题。具体的代码实现细节在此不再赘述,让我们直接查看代码部分。
三、代码实现
#include <iostream>
#include <vector>
#include <algorithm> // for std::sort
using namespace std;
bool backtrack(const vector<int>& nums, int index, int target, vector<int>& size);
bool canPartitionKSubsets(vector<int>& nums, int k) {
int total = 0;
// 统计所有火柴的长度
for (int num : nums) {
total += num;
}
// 如果所有火柴的长度不是k的倍数,直接返回false
if (total == 0 || total % k != 0)
return false;
// 先排序
sort(nums.begin(), nums.end());
// 回溯,从最长的火柴开始
vector<int> size(k, 0); // 初始化长度为k的数组,初始值为0
return backtrack(nums, nums.size() - 1, total / k, size);
}
// index表示访问到当前火柴的位置,target表示正方形的边长,size是长度为k的数组,
// 分别保存正方形k个边的长度
bool backtrack(const vector<int>& nums, int index, int target, vector<int>& size) {
if (index == -1) {
// 如果火柴都访问完了,并且size的k个边的长度都相等,说明是正k边形,直接返回true,
// 否则返回false
for (int i = 1; i < size.size(); i++) {
if (size[i] != size[i - 1]) {
return false;
}
}
return true;
}
// 到这一步说明火柴还没访问完
for (int i = 0; i < size.size(); i++) {
// 如果把当前火柴放到size[i]这个边上,他的长度大于target,我们直接跳过。或者
// size[i] == size[i - 1]即上一个分支的值和当前分支的一样,上一个分支没有成功,
// 说明这个分支也不会成功,直接跳过即可。
if (size[i] + nums[index] > target || (i > 0 && size[i] == size[i - 1]))
continue;
// 如果当前火柴放到size[i]这个边上,长度不大于target,我们就放上面
size[i] += nums[index];
// 然后在放下一个火柴,如果最终能变成正方形,直接返回true
if (backtrack(nums, index - 1, target, size))
return true;
// 如果当前火柴放到size[i]这个边上,最终不能构成正方形,我们就把他从
// size[i]这个边上给移除,然后在试其他的边
size[i] -= nums[index];
}
// 如果不能构成正k边形,直接返回false
return false;
}
int main() {
vector<int> nums = {4, 3, 2, 3, 5, 2, 1};
int k = 4;
bool result = canPartitionKSubsets(nums, k);
if (result) {
cout << "True" << endl;
} else {
cout << "False" << endl;
}
return 0;
}
上面的解法是借鉴了火柴拼正方形的思路,但实际上这道题也可以参考N皇后问题的解决方法。在N皇后问题中,我们首先在第1行的某个位置放置一个皇后,然后是第2行,依此类推,如果能够在第N行都放置皇后,说明找到了一个有效的解。
同样地,这道题也可以采用类似的方法。我们首先计算出第一条边(想象为一个正K边形)的长度为sum/k,然后计算第二条边,以此类推。如果所有边的长度都能达到sum/k,那么我们就可以将数组分割成k个总和相同的非空子集。
#include <iostream>
#include <vector>
#include <algorithm> // for std::sort
using namespace std;
bool dfs(const vector<int>& nums, int total, int start, vector<bool>& visited, int target, int k);
bool canPartitionKSubsets(vector<int>& nums, int k) {
// 统计所有元素的和
int sum = 0;
for (int num : nums)
sum += num;
// 如果不能被k整除,直接返回false
if (sum % k != 0)
return false;
// 因为数组中的每个元素只能使用一次,这里主要用来标记数组中对应的元素是否被使用过
vector<bool> visited(nums.size(), false);
// 先对数组进行排序
sort(nums.begin(), nums.end());
// dfs调用
return dfs(nums, 0, nums.size() - 1, visited, sum / k, k);
}
/**
* 这题把他看做是能不能构成一个正K边形
*
* @param nums 原数组
* @param total 当前边的长度
* @param start 数组中元素开始的位置
* @param visited 数组中对应的元素是否被使用过
* @param target 平均值,也就是每条边的长度
* @param k 计算到第几条边了
* @return
*/
bool dfs(const vector<int>& nums, int total, int start, vector<bool>& visited, int target, int k) {
// 如果k等于0,说明所有的边都计算完了,直接返回true
if (k == 0)
return true;
// 如果当前边的长度是target,说明找到了一条边,继续找剩下的k-1条边,我们看到
// 这里的dfs函数,除了最后k变成了k-1,其他的参数和初始值完全一样
if (total == target)
return dfs(nums, 0, nums.size() - 1, visited, target, k - 1);
// 因为数组是排过序的,我们从大往小找,因为是一条边一条边的查找,所以这里i从start开始,
// start之前的已经被访问过了,不需要在重复访问
for (int i = start; i >= 0; --i) {
// 如果当前元素被使用过,直接跳过
if (visited[i])
continue;
// 如果当前数字太大,也不能选择
if (total + nums[i] > target)
continue;
// 标记为已使用
visited[i] = true;
// 递归,如果成功直接返回true
if (dfs(nums, total + nums[i], i - 1, visited, target, k))
return true;
// 如果不能成功,就把当前值撤销标记
visited[i] = false;
}
return false;
}
int main() {
vector<int> nums = {4, 3, 2, 3, 5, 2, 1};
int k = 4;
bool result = canPartitionKSubsets(nums, k);
if (result) {
cout << "True" << endl;
} else {
cout << "False" << endl;
}
return 0;
}