在问 · 算法题解

698. 划分为k个相等的子集

中等回溯回溯(按桶放置 + 剪枝)力扣原题 ↗

题目描述

输入一个正整数数组 nums 和一个正整数 k。

判断能否把 nums 中的全部元素分到 k 个非空的组里,每个元素恰好属于一组,并且每组元素之和都相同。能则返回 true,否则返回 false。

示例 1
输入:nums = [4,3,2,3,5,2,1], k = 4
输出:true

解释:可以分成 (5)、(1,4)、(2,3)、(2,3),每组和都是 5。

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

解释:总和 10 不能被 3 整除。

数据范围

  • 1 <= k <= nums.length <= 16
  • 0 < nums[i] < 10000
  • 每个值出现的次数在 [1, 4] 之间。

解题思路:回溯(按桶放置 + 剪枝)

每个子集的和必须是 side = sum / k,不能整除就直接失败。把 k 个子集看成 k 个桶,按从大到小的顺序把每个数放进某个桶,放入后桶内和不能超过 side。大数先放,失败会更早暴露。若当前桶与前一个桶的和相同,放进去的结果是对称的,前一个桶已经试过,直接跳过。所有数放完时每个桶恰好等于 side。

为什么正确

进入 dfs(i) 时,前 i 个数已经分配,每个桶的和都不超过 side。

复杂度分析

时间最坏 O(k^n),排序和对称剪枝后实际很快;空间 O(k + n)。

易错点

回溯后忘记把数从桶里拿出来;没有排序和对称剪枝时会超时。还要注意总和不能被 k 整除或最大数超过 side 时直接返回 False。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。

class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        total = sum(nums)
        if total % k:
            return False
        side = total // k
        nums.sort(reverse=True)
        if nums[0] > side:
            return False
        buckets = [0] * k

        def dfs(i):
            if i == len(nums):
                return True
            for b in range(k):
                if buckets[b] + nums[i] > side:
                    continue
                if b > 0 and buckets[b] == buckets[b - 1]:
                    continue
                buckets[b] += nums[i]
                if dfs(i + 1):
                    return True
                buckets[b] -= nums[i]
            return False

        return dfs(0)

载入编辑器,自己跑一遍 →

其他解法

回溯以桶为单位尝试放置,靠剪枝减少搜索;状态压缩把“哪些数已用”当作状态,去掉重复搜索,最坏情况可控。

解法 2:状态压缩动态规划 时间 O(n·2ⁿ),空间 O(2ⁿ)

用二进制掩码 mask 表示已经用掉的数。按顺序填桶,一个桶填满再填下一个,所以当前桶里的和只由 mask 决定,等于已用数之和模 side。dp[mask] 表示这个状态能否合法到达。从可达状态出发,尝试加入一个未用的数,只要不超过 side 就转移。数组先排序,放不下时可以直接结束枚举。全集可达即为真。

每个状态只算一次,最坏复杂度有保证,不依赖剪枝效果。代价是 2ⁿ 的表,只适合 n ≤ 16 这样的小规模。

class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        total = sum(nums)
        if total % k:
            return False
        side = total // k
        nums.sort()
        if nums[-1] > side:
            return False
        n = len(nums)
        full = (1 << n) - 1
        dp = [False] * (1 << n)
        cur = [0] * (1 << n)
        dp[0] = True
        for mask in range(1 << n):
            if not dp[mask]:
                continue
            for i in range(n):
                if cur[mask] + nums[i] > side:
                    break
                if mask >> i & 1:
                    continue
                nxt = mask | 1 << i
                if not dp[nxt]:
                    dp[nxt] = True
                    cur[nxt] = (cur[mask] + nums[i]) % side
        return dp[full]

回溯 专题的其他题目

查看 回溯 全部题目 →