【2.6】回溯算法-解划分为k个相等的子集

一、题目

        给定一个整数数组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;
}