我来用Python3实现LeetCode 2518.好分区的数目。
解题思路
正难则反:直接计算好分区困难,改为计算不好分区(至少一组和<k),再用总分区数减去。
关键点:
1. 如果数组总和 < 2*k,直接返回0
2. 用0-1背包计算所有和 < k 的子集方案数
3. 不好分区数 = 和<k的方案数 × 2(因为第一组或第二组都可能和<k)
4. 好分区数 = 2^n - 不好分区数
Python3代码实现
class Solution:
def countPartitions(self, nums: List[int], k: int) -> int:
MOD = 10**9 + 7
# 如果总和小于2k,不可能有好分区
if sum(nums) < 2 * k:
return 0
n = len(nums)
# dp[j]: 选出若干元素,和为j的方案数 (j < k)
dp = [0] * k
dp[0] = 1 # 空集和为0
# 0-1背包,倒序遍历
for x in nums:
for j in range(k - 1, x - 1, -1):
dp[j] = (dp[j] + dp[j - x]) % MOD
# 计算所有和小于k的方案数
bad = sum(dp) % MOD
# 不好分区数 = 第一组和<k 或 第二组和<k
bad = bad * 2 % MOD
# 总分区数 = 2^n
total = pow(2, n, MOD)
# 好分区数 = 总分区数 - 不好分区数
return (total - bad) % MOD
代码优化版本(更简洁)
class Solution:
def countPartitions(self, nums: List[int], k: int) -> int:
MOD = 10**9 + 7
if sum(nums) < 2 * k:
return 0
dp = [0] * k
dp[0] = 1
for x in nums:
for j in range(k - 1, x - 1, -1):
dp[j] = (dp[j] + dp[j - x]) % MOD
# 总分区数 - 不好分区数
return (pow(2, len(nums), MOD) - sum(dp) * 2) % MOD
复杂度分析
- 时间复杂度:O(n × k),其中n是数组长度,k是给定阈值
- 空间复杂度:O(k),使用一维数组
示例验证
# 示例1
nums = [1,2,3,4], k = 4
# 输出: 6
# 示例2
nums = [3,3,3], k = 4
# 输出: 0
# 示例3
nums = [6,6], k = 2
# 输出: 2
关键点说明
1. 为什么用pow(2, n, MOD):Python内置的快速幂,比循环乘2效率更高
2. 为什么倒序遍历:避免同一个元素被重复使用(0-1背包的标准技巧)
3. 为什么乘以2:第一组和<k 和 第二组和<k 是两种不同的不好分区,且当总和≥2k时不会重叠
4. 边界情况:总和<2k时直接返回0,避免了后续计算的错误