LeetCode 698. 划分为k个相等的子集（回溯算法解决）

Posted 2021-12-11 数据结构和算法

截止到目前我已经写了 600多道算法题，其中部分已经整理成了pdf文档，目前总共有1000多页（并且还会不断的增加），大家可以免费下载
下载链接：https://pan.baidu.com/s/1hjwK0ZeRxYGB8lIkbKuQgQ
提取码：6666

前面我们刚讲过和这题非常类似的题，520，回溯算法解火柴拼正方形，有一点不同的是，第520题判断使用所有的火柴能不能拼成一个正方形。而这题我们可以看做是使用所有的火柴（数组中的元素）能不能拼成一个正K边形。第520题可以看做是一棵四叉树的dfs遍历，那么这题就可以看做是一棵K叉树的dfs遍历，代码非常相似，我们只需要把第520题的代码拿过来修改一下即可，这里就不在过多介绍，直接来看下代码

public boolean canPartitionKSubsets(int[] nums, int k) {
    int total = 0;
    //统计所有火柴的长度
    for (int num : nums) {
        total += num;
    }
    //如果所有火柴的长度不是k的倍数，直接返回false
    if (total == 0 || total % k != 0)
        return false;
    //先排序
    Arrays.sort(nums);
    //回溯，从最长的火柴开始
    return backtrack(nums, nums.length - 1, total / k, new int[k]);
}

//index表示访问到当前火柴的位置，target表示正方形的边长，size是长度为k的数组，
//分别保存正方形k个边的长度
private boolean backtrack(int[] nums, int index, int target, int[] size) {
    if (index == -1) {
        //如果火柴都访问完了，并且size的k个边的长度都相等，说明是正k边形，直接返回true，
        //否则返回false
        // if (size[0] == size[1] && size[1] == size[2] && size[2] == size[3])
        for (int i = 1; i < size.length; i++) {
            if (size[i] != size[i - 1]) {
                return false;
            }
        }
        return true;
    }
    //到这一步说明火柴还没访问完
    for (int i = 0; i < size.length; 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;
}

回溯算法的另一种解法

上面的解法我们是参照火柴拼正方形的解法，其实这题还可以参照N皇后问题来解决，这个可以具体看下394，经典的八皇后问题和N皇后问题和450，什么叫回溯算法，一看就会，一写就废。N皇后是先在第1行的某个位置放一个皇后，然后是第2行……，如果到第N行都能放，说明找到了一个合适的解。

那么这题也一样，我们先求第一条边（想象为正K边形）的长度sum/k，然后再求第二条边……，如果所有的边长度都可以是sum/k，说明我们可以把数组分割成k个总和一样的非空子集。

public boolean canPartitionKSubsets(int[] nums, int k) {
    //统计所有元素的和
    int sum = 0;
    for (int num : nums)
        sum += num;
    //如果不能被4整除，直接返回false
    if (sum % k != 0)
        return false;
    //因为数组中的每个元素只能使用一次，这里主要用来标记数组中对应的元素是否被使用过
    boolean[] visited = new boolean[nums.length];
    //先对数组进行排序
    Arrays.sort(nums);
    //dfs调用
    return dfs(nums, 0, nums.length - 1, visited, sum / k, k);
}

/**
 * 这题把他看做是能不能构成一个正K边形
 *
 * @param nums    原数组
 * @param total   当前边的长度
 * @param start   数组中元素开始的位置
 * @param visited 数组中对应的元素是否被使用过
 * @param target  平均值，也就是每条边的长度
 * @param k       计算到第几条边了
 * @return
 */
public boolean dfs(int[] nums, int total, int start, boolean[] 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.length - 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;
}