DataStructure_Algorithm_HandBook_PreForLeetCode

698. Partition to K Equal Sum Subsets

Given an integer array nums and an integer k, return true if it is 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 is possible to divide it into 4 subsets (5), (1, 4), (2,3), (2,3) with equal sums.

Example 2:

Input: nums = [1,2,3,4], k = 3
Output: false

Constraints:

solution

想象 k=4 的情况,有四个桶,每个元素都有四种选择,但是可以做个剪枝优化,因为每个桶的元素之和不能超过平均值。

按照这个思路就是我们需要开个数组,sums,表示每个桶的元素之和,index 表示当前到哪个位置。

只要我们有一条可行的路径,我们就可以停止了

出口是index 完了之后,看是不是sums 每个元素都等于平均值。

每个元素遍历 k 次

var canPartitionKSubsets = function(nums, k) {
    let sum = 0;
    for(let num of nums)  sum += num;
    // 如果无法被均分,直接返回 false
    if(sum % k !== 0) return false;
    let average = sum / k;
    //对 nums 逆序排序
    nums.sort((a,b)=>b-a);
    return dfs(0, average, new Array(k).fill(0), nums, k);
};

function dfs(index, average, subs, nums, k){
    // 如果所有数字都被分配完毕,检查所有子集是否等于 average
    if(index === nums.length){
        for(let sub of subs) {
            if(sub !== average) return false;
        }
        return true;
    }
    let set = new Set(); 
    // 尝试将当前数字加入每个子集中
    for(let i = 0; i < k; i++){
        // 如果当前子集加上当前数字超过了 average,继续尝试下一个子集
        //我们这边还可以优化一下,因为如果桶的总和是一致的,那么放哪个桶都是一样的组合,因此我们可以加一个 set 数组来优化,也是一种同层去重的思想
        if(nums[index] + subs[i] > average || set.has(subs[i])) continue;
        subs[i] += nums[index];
        // 递归尝试下一个数字
        if(dfs(index + 1, average, subs, nums, k)) return true;
        // 如果加入当前数字后无法得到合法解,则回溯
        subs[i] -= nums[index];
    }
    return false; 
}

上面还有一个思路,就是先把大的元素排序,这样可以更快的得到结果。 所以就可以对nums 进行逆序排序。

当然力扣上面这道题目是超时的,因此还是后面学动态规划的时候用更优秀的解法。

这个遍历子集的方法可以通过

var canPartitionKSubsets = function(nums, k) {
    let sum = 0;
    for (let num of nums) sum += num;
    if (sum % k !== 0) return false;
    let targetSum = sum / k;
    let visited = new Array(nums.length).fill(false);
    return backtrack(nums, k, targetSum, 0, 0, visited);
};

function backtrack(nums, k, targetSum, currentSum, start, visited) {
    if (k === 0) return true; // 所有子集已经满足条件
    if (currentSum === targetSum) return backtrack(nums, k - 1, targetSum, 0, 0, visited); // 开始下一个子集
    for (let i = start; i < nums.length; i++) {
        if (!visited[i] && currentSum + nums[i] <= targetSum) {
            visited[i] = true;
            if (backtrack(nums, k, targetSum, currentSum + nums[i], i + 1, visited)) return true;
            visited[i] = false; // 回溯
        }
    }
    return false;
}

backtrack 函数负责生成符合条件的子集。它使用了一个额外的数组 visited 来记录哪些元素已经被访问过,避免重复。在每次递归调用中,我们都尝试将当前元素添加到子集中,然后继续递归地向下搜索。如果当前子集满足条件,则尝试开始下一个子集;如果遍历完所有元素后仍然无法满足条件,则回溯到上一个状态。