LeetCode - 698 划分为k个相等的子集

目录

题目来源

题目描述

示例

提示

题目解析

算法源码

---

题目来源

698. 划分为k个相等的子集 - 力扣（LeetCode）

题目描述

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

示例

提示

• 1 <= k <= len(nums) <= 16
• 0 < nums[i] < 10000
• 每个元素的频率在 [1,4] 范围内

题目解析

本题其实是一道子集问题，可以利用回溯法求解。

首先，我们可以求出nums数组的所有元素之和sum，如果sum % k 不为0的话，说明每一份子集的和不是整数，但是题目说了nums数组中的元素都是整数，因此每一个子集的和也应该是整数，因此如果sum % k 不为0的话，则返回false。

如果sum % k 为0，则我们假设subSum = sum / k，即每一个子集的和为subSum，那么subSum应该大于等于nums数组的最大值，因为nums数组的每一个子集至少有一个元素组成，因此subSum应该大于或等于每一个nums元素，即至少不小于nums数组的最大值。

当面两个条件都符合后，还不能说明nums可以被分为k个和相等的子集，比如nums=[2,2,2,2,3,4,5] ，且k=4。

nums数组的和sum = 20，k=4，因此subSum = 5,而subSum >= max(nums[i]) = 5，因此符合上面两个条件，但是nums数组却无法分为4个等和子集。

因此，我们需要一种算法来判断一个数组是否可以划分k个和相等的子集，此时我们就需要借助回溯算法。

我们可以想象划分子集问题，可以看成向k个桶中放球，每个桶的承重为subSum，而现在nums中存放着重量不一的多个球，如果nums中的球都能放到k个桶中，则可以划分子集成功。

比如nums = [4, 3, 2, 3, 5, 2, 1], k = 4，对应的subSum=5

首先，nums的重量5的球可以独占一个桶，我们可以减少考虑一个桶的处理

 接下来，我们需要将球按照降序排序，然后依次放入桶中

 

由于 4 + 3 > 5，因此球3需要放到下一个桶中

同理处理下面的球

这里我们对nums进行了降序排序，而不是升序排序，因为升序排序可能会产生非常复杂的放置行为，这将不利于我们理解。 

上述逻辑对应的算法代码如下

JS

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {boolean}
 */
var canPartitionKSubsets = function(nums, k) {
  const sum = nums.sort((a,b) => b - a).reduce((p, c) => p + c)
 
  if(sum % k !== 0) return false
 
  const subSum = sum / k
 
  if(subSum < nums[0]) return false
 
  while(nums[0] === subSum) {
      nums.shift()
      k--
  }
 
  const buckets = new Array(k).fill(0)
 
  return partition(nums, 0, buckets, subSum)
};
 
function partition(nums, index, buckets, target) {
    if(index === nums.length) return true
 
    const selected = nums[index]
 
    for(let i = 0; i < buckets.length; i++) {
        if(selected + buckets[i] <= target) {
            buckets[i] += selected
            if(partition(nums, index+1, buckets, target)) return true
            buckets[i] -= selected
        }
    }
    return false
}

可以发现，性能非常差。

有这样一种情况，即我们选择了一个球，但是发现此时每一个桶中的球重量之和都相等，也就是说此时选择的球放到哪一个桶中，最终结果都是相同的，比如球放到桶1中重量超了，则放到其他桶中重量肯定也是超的。

但是根本上面算法逻辑，我们肯定会在选择的球放到桶1超了之后，继续到桶2中尝试，而桶2超了之后，必然放到桶3中尝试，这其实就是一种无意义的试探操作。

因此我们再进入递归流程前，先判断 buckets[i]的重量是否和buckets[i-1]的重量是否相等，若相等的话，则buckets[i-1]试过了，则buckets[i]就不用试了。

最终代码如下

JavaScript算法源码

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {boolean}
 */
var canPartitionKSubsets = function(nums, k) {
  const sum = nums.sort((a,b) => b - a).reduce((p, c) => p + c)
 
  if(sum % k !== 0) return false
 
  const subSum = sum / k
 
  if(subSum < nums[0]) return false
 
  while(nums[0] === subSum) {
      nums.shift()
      k--
  }
 
  const buckets = new Array(k).fill(0)
 
  return partition(nums, 0, buckets, subSum)
};
 
function partition(nums, index, buckets, target) {
    if(index === nums.length) return true
 
    const selected = nums[index]
 
    for(let i = 0; i < buckets.length; i++) {
        if (i > 0 && buckets[i] === buckets[i - 1]) continue // 此步剪枝将大幅提升性能
        if(selected + buckets[i] <= target) {
            buckets[i] += selected
            if(partition(nums, index+1, buckets, target)) return true
            buckets[i] -= selected
        }
    }
    return false
}

Java算法源码

import java.util.Arrays;
 
class Solution {
  public boolean canPartitionKSubsets(int[] nums, int k) {
    int sum = 0;
    for (int num : nums) sum += num;
    if (sum % k != 0) return false;
 
    int subSum = sum / k;
 
    // nums升序
    Arrays.sort(nums);
 
    // nums元素反转，实则实现降序
    int l = 0;
    int r = nums.length - 1;
    while (l < r) {
      int tmp = nums[l];
      nums[l] = nums[r];
      nums[r] = tmp;
      l++;
      r--;
    }
 
    if (subSum < nums[0]) return false;
 
    int[] buckets = new int[k];
    return partition(nums, 0, buckets, subSum);
  }
 
  public boolean partition(int[] nums, int index, int[] buckets, int target) {
    if (index == nums.length) return true;
 
    int selected = nums[index];
    for (int i = 0; i < buckets.length; i++) {
      if (i > 0 && buckets[i] == buckets[i - 1]) continue;
      if (selected + buckets[i] <= target) {
        buckets[i] += selected;
        if (partition(nums, index + 1, buckets, target)) return true;
        buckets[i] -= selected;
      }
    }
 
    return false;
  }
}

Python算法源码

class Solution(object):
    def canPartitionKSubsets(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: bool
        """
        total = sum(nums)
        if total % k != 0:
            return False
        
        avg = total // k
 
        nums.sort(reverse=True)
        if avg < nums[0]:
            return False
        
        buckets = [0]*k
        return partition(nums, 0, buckets, avg)
    
def partition(nums, index, buckets, target):
    if index == len(nums):
        return True
    
    selected = nums[index]
    for i in range(len(buckets)):
        if i>0 and buckets[i] == buckets[i-1]:
            continue
        if selected + buckets[i] <= target:
            buckets[i] += selected
            if partition(nums, index+1, buckets, target):
                return True
            buckets[i] -= selected
    return False