LeetCode #698 Partition to K Equal Sum Subsets

题目

Given an array of integers nums and a positive integer k, find whether it’s possible to divide this array into k non-empty subsets whose sums are all equal.

Example 1:

Input: nums = [4, 3, 2, 3, 5, 2, 1], k = 4
Output: True
Explanation: It's possible to divide it into 4 subsets (5), (1, 4), (2,3), (2,3) with equal sums.

Note:

1 <= k <= len(nums) <= 16

.

0 < nums[i] < 10000

.

解题思路

首先确定

k

个子集的相同的和

target

是什么。方法是把数组

nums

中的所有数字求和得到

sum

，如果

sum % k != 0

，则肯定不能划分为

k

个相同和的子集；否则，

target = sum / k

。

采用 DFS 的方法，递归计算每一个子集：每次从数组中取一个没有被访问过的数（即这个数没有被归入到任何一个子集中去），加入到当前递归层计算的子集中，如果加入这个数后当前子集的和与

target

相等，则重新从数组第一个元素开始，选择没有被访问过的数计算下一个子集，并设置

k = k - 1

表示已经找到了一个子集，还剩下

k - 1

个子集需要寻找；如果加入这个数后当前子集的和已经超过

target

，则

return false

，并在返回上一层递归时重新将这个数设置为没有被访问；如果加入这个数后当前子集的和仍然小于

target

，则循环这个数后面的数继续添加进当前子集中并进行判断。

当

k == 1

时，说明前面的数已经被划分成了求和相同的

k - 1

个子集，则剩下的数一定能组成第

k

个这样的子集（因为剩下数字的和必然等于

sum - [target * (k - 1)] = target

），此时说明分组成功，则

return true

。

C++代码实现

class Solution {
public:
bool canPartitionKSubsets(vector<int>& nums, int k) {
int sum = 0;
for (int i = 0; i < nums.size(); ++i) { sum += nums[i]; }
if (sum % k != 0) { return false; }

vector<bool> visited(nums.size(), false);
return canPartition(nums, visited, 0, 0, k, sum / k);
}

// startIndex 表示从第几个位置开始选后面的数加入当前子集中
// curSum 表示当前子集的所有元素的和
// k 为剩下的子集数
bool canPartition(vector<int>& nums, vector<bool>& visited, int startIndex, int curSum, int k, int target) {
if (k == 1) { return true; }
if (curSum == target) { return canPartition(nums, visited, 0, 0, k - 1, target); }
if (curSum > target) { return false; }

for (int i = startIndex; i < nums.size(); ++i) {
if (!visited[i]) {
visited[i] = true;
if (canPartition(nums, visited, i + 1, curSum + nums[i], k, target)) { return true; }
visited[i] = false;
}
}

return false;
}
};

拓展

题目中限定了

0 < nums[i] < 10000

，当

nums[i]

中存在负数时，需要对代码进行一定的改动：

计算当前递归层的子集时，如果加入的数使得当前子集的和超过了

target

，此时不能直接

return false

，因为后面可能加入负数后可以使得当前子集的和等于

target

。

当

target == sum / k == 0

时，注意到空集的和也是

0

，此时要加入一个计数器

elemNum

计算当前子集的元素个数，防止出现空集。

C++ 代码实现如下：

class Solution {
public:
bool canPartitionKSubsets(vector<int>& nums, int k) {
int sum = 0;
for (int i = 0; i < nums.size(); ++i) { sum += nums[i]; }
if (sum % k != 0) { return false; }

vector<bool> visited(nums.size(), false);
return canPartition(nums, visited, 0, 0, 0, k, sum / k);
}

// startIndex 表示从第几个位置开始选后面的数加入当前子集中
// curSum 表示当前子集的所有元素的和
// elemNum 表示当前子集的元素个数
// k 为剩下的子集数
bool canPartition(vector<int>& nums, vector<bool>& visited,
int startIndex, int curSum, int elemNum, int k, int target) {
if (k == 1) { return true; }
4000
if (curSum == target && elemNum > 0) {
return canPartition(nums, visited, 0, 0, 0, k - 1, target);
}

for (int i = startIndex; i < nums.size(); ++i) {
if (!visited[i]) {
visited[i] = true;
if (canPartition(nums, visited, i + 1, curSum + nums[i], elemNum + 1, k, target)) {
return true;
}
visited[i] = false;
}
}

return false;
}
};