[数据结构与算法]Leetcode 698. 划分为k个相等的子集

1.题目描述

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

输入： nums = [4, 3, 2, 3, 5, 2, 1], k = 4
输出： True
说明： 有可能将其分成 4 个子集（5），（1,4），（2,3），（2,3）等于总和。

输入: nums = [1,2,3,4], k = 3
输出: false

提示：

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

2.思路分析

特殊情况判断:

• 如果数组 nums 的和 sum 不能整除 k，或者数组 nums 中的最大值大于每个子集的和 sum / k (也就是 target)，返回 false。

totalSum = sum(nums)
# 预处理 特殊情况处理
if totalSum % k != 0:
    return False
target = totalSum // k
bucket = [0] * (k + 1)
# 排序(降序) 方便后面剪枝
nums.sort(reverse=True)
# 如果最小的元素大于target 直接返回False
if nums[-1] > target:
    return False

以数字为视角: 每个数字每一次的可选择项都为「所有桶」，所以每一次可选择项均为桶的数量 k。故时间复杂度指数的底数为 k

回溯三部曲:

确定回溯函数要处理的参数以及返回值

• **参数:**num: 输入数组 k: 桶个数 startIndex: 当前元素对应索引 target: 目标值
• 返回值:boolean

def backtrack(nums: List[int], k: int, bucket: List[int], startIndex: int, target: int) -> bool:

确定终止条件

nums 中的数字全部消耗完，说明找到了 k 个桶，且每个桶内的数字和为 target。

if startIndex == len(nums):
    # 判断现在桶中的元素是否符合要求 -> 路径是否满足要求
	for i in range(k):
		if bucket[i] != target:
			return False
	return True

单层搜索逻辑

• 依次遍历每个桶，尝试将当前的数字 nums[startIndex] 放入到该桶中，看最终是否能够凑成 k 个桶。
• 如果当前遍历到的桶加上 nums[startIndex] > 目标值 target，则说明 nums[startIndex] 不应该放到当前桶内，继续遍历下一个桶。
• 当前桶的数字和与上一个桶一样，跳过。

# 单层搜索逻辑
for i in range(k):
    # 去重优化
    if i > 0 and bucket[i] == bucket[i - 1]:
        continue
    # 剪枝: 如果当前元素的值+bucket[i]>target,选择下一个桶
    if bucket[i] + nums[startIndex] > target:
        continue
    # 处理元素
    bucket[i] += nums[startIndex]
    # 处理下一个元素(递归)
    if backtrack(nums, k, bucket, startIndex + 1, target):
        return True
    # 撤销选择
    bucket[i] -= nums[startIndex]
# 遍历结束,k个桶都不满足要求
return False

3.代码实现

class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        totalSum = sum(nums)
        # 预处理 特殊情况处理
        if totalSum % k != 0:
            return False
        target = totalSum // k
        bucket = [0] * (k + 1)
        # 排序优化
        nums.sort(reverse=True)
        if nums[-1] > target:
            return False

        # 回溯函数要处理的参数以及返回值
        def backtrack(nums: List[int], k: int, bucket: List[int], startIndex: int, target: int) -> bool:
            """
            Args:
            :param nums: 输入数组
            :param k: 桶个数
            :param startIndex: 当前元素对应索引
            :param target: 目标值
            :return:
            """
            # 回溯终止条件:处理完所有元素以及当前桶中的元素是否符合要求
            # 结束条件优化
            if startIndex == len(nums):
                # 判断现在桶中的元素是否符合要求 -> 路径是否满足要求
                for i in range(k):
                    if bucket[i] != target:
                        return False
                return True
            # 单层搜索逻辑
            for i in range(k):
                # 去重优化
                if i > 0 and bucket[i] == bucket[i - 1]:
                    continue
                # 剪枝: 如果当前元素的值+bucket[i]>target,选择下一个桶
                if bucket[i] + nums[startIndex] > target:
                    continue
                # 处理元素
                bucket[i] += nums[startIndex]
                # 处理下一个元素(递归)
                if backtrack(nums, k, bucket, startIndex + 1, target):
                    return True
                # 撤销选择
                bucket[i] -= nums[startIndex]
            # 遍历结束,k个桶都不满足要求
            return False
        return backtrack(nums, k, bucket, 0, target)

复杂度分析: todo