集合划分问题：思路与剪枝[回溯]

集合划分问题[回溯]

集合划分问题：将数组划分成K个非空子集，每个子集都满足指定条件，比如元素和为指定值。

一、回溯思想

题目：LeetCode 698. 划分为k个相等的子集

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

由于划分成K个不同的非空子集，因此我们不妨使用一个数组记录K个子集的元素和分别是多少。

由于数组nums的总和确定为sum，因此K个非空子集中，每个子集的和应当为指定值sum/K。

选择列表

对于第idx个元素，我们可以选择将其放入K个子集中的任意一个。如果放入后，该子集的和不超出指定值，则递归至下一层。

递归

如果第idx个元素已放入，则递归到下一层，选择将第idx+1个元素放入某一个子集。

回溯

如果下一层的递归结果返回false，则说明第idx个元素不能放入该集合。

此时需要将第idx个元素从先前选择的集合中取出，同时在选择列表中重新选择一个集合。

结束条件

当idx==n时，说明所有元素都已经放入K个子集中。且可以证明：每个子集的和都是指定值。此时返回true。

根据上述分析，可以写出这样的代码：

class Solution {
public:
    bool dfs(vector<int>& nums, int idx, vector<int>& v, int target)
    {
        // 结束条件：n个元素都已放入
        if (idx == nums.size())
            return true;
        for (int i = 0; i < v.size(); ++i)
        {
            // v[i] + nums[idx] > target，说明第idx个球放入会导致集合的和超出目标值
            if (v[i] + nums[idx] > target)
                continue;
            
            v[i] += nums[idx];
            // 第idx个数放入第i个集合后，和不会超出target，则继续向下搜索
            if (dfs(nums, idx + 1, v, target) == true)
                return true;
            // 回溯，将第idx个元素取出，重新放入新的集合
            v[i] -= nums[idx];
        }
        // k个集合都放入失败，返回false
        return false;
    }
    bool canPartitionKSubsets(vector<int>& nums, int k) 
    {
        int n = nums.size();
        int sum = 0;
        for (auto e : nums)
        {
            sum += e;
        }

        // 总和无法被k整除，必然不可以划分
        if (sum % k != 0)
            return false;

        // 每个集合的总和都是sum/k
        int target = sum / k;
        
        vector<int> v(k);
        return dfs(nums, 0, v, target);
    }
};

显然，没有剪枝与优化的暴力回溯会超时。

二、优化与剪枝

• 首先，我们为原数组进行从大到小的排序。

  分析一下这样做的好处：

  当数组从大到小排序时，我们优先将大的元素放入集合。

  1. 在极端条件下，如果最大的元素超出了指定值，那么该元素显然无法放入任意一个集合，此时递归直接就结束了。
  2. 在一般条件下，当一个集合已经放入一个较大的数时，其它可以选择放入该集合的数就比较少了，这样可以减少递归的次数。如果我们选择从小到大将数放入集合，那么一个小的数放入集合后，该集合还可以放入更多小的数，明显增加了递归的深度和次数。

• 其次，我们考虑如何剪枝。

  当我们尝试将第idx个球放入第i个集合时，说明将它放入前i-1个集合都已经失败了。那么此时如果第i个集合的元素和与第i-1个集合的元素和相等，则说明第idx个球再放入第i个集合，依然会失败。因此该递归分支可以被剪掉。

优化后的代码：

class Solution {
public:
    bool dfs(vector<int>& nums, int idx, vector<int>& v, int target)
    {
        if (idx == nums.size())
            return true;
        for (int i = 0; i < v.size(); ++i)
        {
            // v[i] + nums[idx] > target，说明第idx个球放入会导致集合的和超出目标值
            // v[i - 1] == v[i]，且遍历到第i个集合时说明第idx个球放入第i-1个已经失败了，那么再放入第i个依然会失败，因此将这个递归分支剪掉
            if (v[i] + nums[idx] > target || (i > 0 && v[i - 1] == v[i]))
                continue;
            v[i] += nums[idx];
            // 第idx个数放入第i个集合后，和不会超出target，则继续向下搜索
            if (dfs(nums, idx + 1, v, target) == true)
                return true;
            v[i] -= nums[idx];
        }
        return false;
    }
    bool canPartitionKSubsets(vector<int>& nums, int k) 
    {
        int n = nums.size();
        int sum = 0;
        for (auto e : nums)
        {
            sum += e;
        }

        // 总和无法被k整除，必然不可以划分
        if (sum % k != 0)
            return false;

        // 每个集合的总和都是sum/k
        int target = sum / k;
        
        vector<int> v(k);
        sort(nums.begin(), nums.end(), greater<int>());
        return dfs(nums, 0, v, target);
    }
};

三、变种题型

LeetCode 2305. 公平分发饼干

class Solution {
public:
    bool dfs(vector<int>& cookies, vector<int>& v, int idx, int target)
    {
        if (idx == cookies.size())
            return true;
        for (int i = 0; i < v.size(); ++i)
        {
            if (v[i] + cookies[idx] > target || (i && v[i - 1] == v[i]))
                continue;
            v[i] += cookies[idx];
            if (dfs(cookies, v, idx + 1, target))
                return true;
            v[i] -= cookies[idx];
        }
        return false;
    }
    int distributeCookies(vector<int>& cookies, int k) 
    {
        sort(cookies.begin(), cookies.end(), greater<int>());
        // 二分查找不公平程度target
        // 将n个元素分成k个集合，集合的元素和不能超出target
        // target的取值范围为[min, sum/k]
        // 其中min为cookies中的最小值，sum为cookies的和
        int sum = accumulate(cookies.begin(), cookies.end(), 0);
        int l = cookies.back();
        int r = sum;
        int ans = 0;
        
        while (l <= r)
        {
            int mid = l + ((r - l) >> 1);

            vector<int> v(k);
            if (dfs(cookies, v, 0, mid))
            {
                ans = mid;
                r = mid - 1;
            }
            else
            {
                l = mid + 1;
            }
        }
        return ans;
    }
};