Code & Func

2022 年新一轮的打卡，准备每天都完成LeetCode的每日一题，并写下题解。今天的题目是698. Partition to k equal sum subsets

给定一个整数数组 nums 和一个正整数 k，找出是否有可能把这个数组分成 k 个非空子集，其总和都相等。

示例 1：

输入： nums = [4, 3, 2, 3, 5, 2, 1], k = 4 输出： True 说明：有可能将其分成 4 个子集（5），（1,4），（2,3），（2,3）等于总和。

示例 2:

输入: nums = [1,2,3,4], k = 3 输出: false

提示：

1 <= k <= len(nums) <= 16
0 < nums[i] < 10000
每个元素的频率在 [1,4] 范围内

题目的意思很简单，就是判断是否能把数组中元素划分为 k 组，每组的元素之和相等。因此很容易得到一个快速判断的条件：当数组元素之和不能整除 k 时，则不可能划分。如果能整除，则 m = sum(nums) / k为每组的元素之和。

从提示来看，这个算法的复杂度估计很高，因为它给的输入范围就很小，数组元素个数最大才到16。

由于给的输入数据不大，所以可以考虑直接通过回溯暴力计算出来，是否能划分为k组。

解法 1：回溯法

回溯时可以先准备 k 个大小为 m 的桶，然后每次对数组中的一个元素，尝试将其放入某个桶中（需要保证桶内大小大于等于该元素大小），如果能把所有元素都放入桶内，则可以划分为 k 组。同时可以在进行回溯前对数组进行排序，然后按从大到小的顺序尝试将元素放入桶中。

1
bool canPartitionKSubsets(vector<int>& nums, int k) {
2
  int sum = 0;
3
  for(auto i : nums) {
4
    sum += i;
5
    imap[i]++;
6
  }
7

8
  if (sum % k) return false;
9
  int m = sum / k;
10

11
  vector<int> buckets(k, m);
12
  sort(nums.begin(), nums.end());
13

14
  return dfs(nums, buckets, nums.size() - 1);
15
}
14 collapsed lines
16

17
bool dfs(vector<int>& nums, vector<int>& buckets, int index) {
18
  if (index < 0) return true;
19
  for (int i = 0;i < buckets.size();i++) {
20
    if (nums[index] <= buckets[i]) {
21
      buckets[i] -= nums[index];
22
      if (dfs(nums, buckets, index-1)) {
23
        return true;
24
      }
25
      buckets[i] += nums[index];
26
    }
27
  }
28
  return false;
29
}

上述算法还是没法直接AC的，需要加入各种剪枝条件。

解法2：回溯+剪枝

剪枝条件：

放入新的元素后，桶内剩余容积要大于最小元素大小：buckets[i] - nums[index] > nums[0]
如果两个桶容积相同，则放哪个效果都一样，buckets[i] == buckets[i - 1]
将要放入桶内的第 i 个元素，放在第 i - 1 个桶之后的哪个桶都是一样的，nums.size() < i + index

加入剪枝条件后的代码：

1
bool canPartitionKSubsets(vector<int>& nums, int k) {
2
  int sum = 0;
3
  unordered_map<int, int> imap;
4
  for(auto i : nums) {
5
    sum += i;
6
    imap[i]++;
7
  }
8

9
  if (sum % k) return false;
10
  int m = sum / k;
11

12
  vector<int> buckets(k, m);
13
  sort(nums.begin(), nums.end());
14

15
  return dfs(nums, buckets, nums.size() - 1);
17 collapsed lines
16
}
17

18
bool dfs(vector<int>& nums, vector<int>& buckets, int index) {
19
  if (index < 0) return true;
20
  for (int i = 0;i < buckets.size();i++) {
21
    if (nums.size() < i + index) break;
22
    if (i > 0 && buckets[i] == buckets[i - 1]) continue;
23
    if (nums[index] == buckets[i] || buckets[i] - nums[index] >= nums[0]) {
24
      buckets[i] -= nums[index];
25
      if (dfs(nums, buckets, index-1)) {
26
        return true;
27
      }
28
      buckets[i] += nums[index];
29
    }
30
  }
31
  return false;
32
}

解法3 状态压缩 +动态规划

后续学习学习

官方题解