[模版总结] - 集合划分类DFS模版

题目描述

给定一个数组，给定一个数字k, 问能不能讲数组内数等分成k份，使得每一个集合中数和相等。

题目链接

下面两道题问题及其类似，可作为同一类题目思考

Leetcode 698

Leetcode 473

题目思路

这道题是一道经典集合划分类问题，这类问题问询通常是判读是否将一个集合里面的所有元素划分到不同满足一定条件的子集合中。

集合划分也是思路和排列组合类似，可以想象成给球分桶的问题，这种类型的题目通常有两个思考方式：给每一个球选桶，给每一个桶选球

1. 给每一个球选桶：我们每一层递归的是球，每一个球每层递归可以选择进入一个桶，以此类推。

2. 给每一个桶选球：我们每一层递归的是桶，相当于每一层给某一个桶塞球，当塞满了就开始进入下一个桶。

这里的球就是指集合中的每一个数，而桶则是指等分的每一个子集合。

1. 暴力DFS

给每一个数选择一个子集合(超时)- 思路很简单，每一层递归是给当前数字选择一个子集合，如果当前子集合超过等分值那么跳过，直到所有数字选完子集合后，比较子集合中是否相等即可。

代码如下：

class Solution {
    int[] arr;
    int avg;
    public boolean canPartitionKSubsets(int[] nums, int k) {
        int sum = 0;
        for (int i: nums) sum+=i;

        avg = sum/k;
        if (avg*k!=sum) return false;
        arr = nums;

        return dfs(0, new int[k]);

    }

    private boolean dfs(int idx, int[] bkts) {
        if (idx==arr.length) {
            // 所有球找完桶，这时确定每个桶里面的数字
            for (int i: bkts) {
                if (i!=avg) return false;
            }
            return true;
        }

        // 开始选桶
        for (int i=0; i<bkts.length; i++) {
            if (bkts[i]+arr[idx] > avg) continue;

            // 桶里还能加
            bkts[i]+=arr[idx];
            if (dfs(idx+1, bkts)) return true;
            bkts[i]-=arr[idx]; // 组合类题目要回溯

        }

        return false;
    }
}

时间复杂度：, 每一个数字有k个选择方式；空间复杂度：，如果忽略递归占用的额外占空间。

给每一个子集合加入数(通过) - 相当于给每一个子集合选择一套数字，如果这个子集合塞满了，则继续下一个子集合直到所有的“桶”都塞完了“球”。思路和球选桶类似，但是需要开额外空间加入标记位来确认当前球被选过没有。

代码如下：

class Solution {
    int[] arr;
    int avg;
    public boolean canPartitionKSubsets(int[] nums, int k) {
        int sum = 0;
        for (int i: nums) sum+=i;

        avg = sum/k;
        if (avg*k!=sum) return false;
        arr = nums;

        return dfs(0, 0, new int[k], new boolean[nums.length]);
    }

    private boolean dfs(int idx, int bkt, int[] bkts, boolean[] visited) {
        if (bkt==bkts.length) {
            for (int i: bkts) {
                if (i!=avg) return false;
            }
            return true;
        }

        if (bkts[bkt]==avg) {
            // 目前选满了，换下一个桶
            return dfs(0, bkt+1, bkts, visited);
        }


        // 当前桶塞球
        for (int i=idx; i<arr.length; i++) {
            if (visited[i]) continue; //球用过了
            if (bkts[bkt]+arr[i]>avg) continue;

            // 当前桶可以塞
            visited[i] = true;
            bkts[bkt]+=arr[i];
            if (dfs(i+1, bkt, bkts, visited)) return true;
            visited[i] = false;
            bkts[bkt]-=arr[i]; // 组合问题回溯
            
        }
        return false;
    }
}

时间复杂度：, 比球选桶的思路快了不少, 每一个桶在选球时有个选法；空间复杂度：，需要开额外空间进行去重。

2. DFS + 剪枝优化

给每一个数选择一个子集合（通过） - 不难理解，在递归寻找解的过程中，有很多重复解的情况，比如桶1选择2，3，4号球，桶2选择5，6号球，这种情况和桶1选择5，6号球，桶2选择2，3，4号球是重复情况。我们只是需要知道组合数而不需要知道排序。

代码如下:

class Solution {
    int[] arr;
    int avg;
    public boolean canPartitionKSubsets(int[] nums, int k) {
        int sum = 0;
        for (int i: nums) sum+=i;

        avg = sum/k;
        if (avg*k!=sum) return false;
        arr = nums;

        return dfs(0, new int[k]);

    }

    private boolean dfs(int idx, int[] bkts) {
        if (idx==arr.length) {
            // 所有球找完桶，这时确定每个桶里面的数字
            for (int i: bkts) {
                if (i!=avg) return false;
            }
            return true;
        }

        // 开始选桶
        for (int i=0; i<bkts.length; i++) {
            if (bkts[i]+arr[idx] > avg) continue;
            
            // 相邻两个桶值相同，那么选了上一个就不需要选这个
            if (i>0 && bkts[i-1]==bkts[i]) continue;

            // 桶里还能加
            bkts[i]+=arr[idx];
            if (dfs(idx+1, bkts)) return true;
            bkts[i]-=arr[idx]; // 组合类题目要回溯

        }

        return false;
    }
}

时间复杂度：最坏, 具体复杂度很难估计；空间复杂度：，如果忽略递归占用的额外占空间。

给每一个子集合加入数(通过) - 给桶塞球的思路我们可能会给两个桶塞入相同组合的球，那么可以设置一个HashMap记录之前的球选择，如果球组合被选过那么直接返回结果而不在往后计算。

注：这里还有一个优化是对boolean[] visited数组进行空间优化，我们可以直接用一个N位 int 来保存访问信息，但你需要熟悉下面几个位操作 - 非常巧妙

1. ((int >> i) & 1) == 1: 判断第i个球是否用过 等价于 visited[i]==true

2. int |= 1<<i: 给第i个球标记为访问 等价于 visited[i] = true

3. int ^= 1<<i: 给第i个球回溯为未访问 等价于 visited[i] = false

代码如下：

class Solution {
    int[] arr;
    int limit;
    HashMap<Integer, Boolean> paths = new HashMap<>();
    public boolean canPartitionKSubsets(int[] nums, int k) {
        int sum = 0;
        for (int i: nums) sum+=i;
        boolean[] visited = new boolean[nums.length];

        int avg = sum/k;
        if (avg*k!=sum) return false;

        arr = nums;
        limit = avg;

        // 骚套路剪枝通过int来记录每一个元素使用与否
        int used = 0;

        return dfs(used, 0, new int[k], 0);
    }

    private boolean dfs(int used, int bkt, int[] bkts, int start) {
        // 以k个桶的角度选择候选
        if (bkt==bkts.length) return true;
        if (bkts[bkt]==limit) {
            boolean res =  dfs(used, bkt+1, bkts, 0);
            paths.put(used, res);
            return res;
        } 

        if (paths.containsKey(used)) {
            return paths.get(used);
        }

        for (int i=start; i<arr.length; i++) {
            if (((used>>i) & 1)==1) continue;
            if (bkts[bkt]+arr[i]>limit) continue;
            

            used |= 1 << i;
            bkts[bkt]+=arr[i];
            if (dfs(used, bkt, bkts, start+1)) return true;
            bkts[bkt]-=arr[i];
            used ^= 1 << i;
        }

        return false;
    }
}

时间复杂度：最坏, 比球选桶的思路快了不少, 每一个桶在选球时有个选法；空间复杂度：，虽然省去了visited的空间，但是需要开额外空间进行路径去重。