和带限制的子多重集合的数目

标签: 数组 哈希表 动态规划 滑动窗口

难度: Hard

给你一个下标从 0 开始的非负整数数组 nums 和两个整数 l 和 r 。

请你返回 nums 中子多重集合的和在闭区间 [l, r] 之间的 子多重集合的数目

由于答案可能很大,请你将答案对 109 + 7 取余后返回。

子多重集合 指的是从数组中选出一些元素构成的 无序 集合,每个元素 x 出现的次数可以是 0, 1, ..., occ[x] 次,其中 occ[x] 是元素 x 在数组中的出现次数。

注意:

  • 如果两个子多重集合中的元素排序后一模一样,那么它们两个是相同的 子多重集合 。
  •  集合的和是 0 。

示例 1:

输入:nums = [1,2,2,3], l = 6, r = 6
输出:1
解释:唯一和为 6 的子集合是 {1, 2, 3} 。

示例 2:

输入:nums = [2,1,4,2,7], l = 1, r = 5
输出:7
解释:和在闭区间 [1, 5] 之间的子多重集合为 {1} ,{2} ,{4} ,{2, 2} ,{1, 2} ,{1, 4} 和 {1, 2, 2} 。

示例 3:

输入:nums = [1,2,1,3,5,2], l = 3, r = 5
输出:9
解释:和在闭区间 [3, 5] 之间的子多重集合为 {3} ,{5} ,{1, 2} ,{1, 3} ,{2, 2} ,{2, 3} ,{1, 1, 2} ,{1, 1, 3} 和 {1, 2, 2} 。

提示:

  • 1 <= nums.length <= 2 * 104
  • 0 <= nums[i] <= 2 * 104
  • nums 的和不超过 2 * 104
  • 0 <= l <= r <= 2 * 104

Submission

运行时间: 176 ms

内存: 17.4 MB

class Solution:
    def countSubMultisets(self, nums: List[int], l: int, r: int) -> int:
        s = sum(nums)
        if s < l:
            return 0
        r, S = min(r, s), -1
        c, m = Counter(nums), 10**9 + 7
        k = s + 1 >> 1
        if l >= k:
            l, r = s - r, s - l
        elif r >= k:
            l, r = sorted((l - 1, s - r - 1))
            S = reduce(lambda i, j: i * j % m, (v + 1 for v in c.values()))
            if r < 0:
                return S
        s, r1 = 1, r + 1
        a = [0] * r1
        a[0] = c[0] + 1
        del c[0]
        for v, k in sorted((k * v, k) for k, v in c.items()):
            s = min(s + v, r1)
            a1 = a.copy()
            v += k
            for i in range(k, s):
                a[i] += a[i - k]
                if (j := i - v) >= 0:
                    a[i] -= a1[j]
                a[i] %= m
        if S < 0:
            return sum(a[l:]) % m
        return (S - sum(a[:l + 1]) - sum(a)) % m

Explain

此题解使用动态规划的方法解决问题。首先,它计算输入数组的总和 s,如果 s 小于 l,则直接返回 0。接着,它使用 Counter 来统计数组中每个元素的出现次数。动态规划数组 a 用于存储从和为 0 到 r 的子集的计数。a 的初始状态是 a[0] 等于 0 的元素的数量加 1。对于每个元素和其计数,更新动态规划数组 a,考虑包含不同数量的当前元素时的子集和。最终,根据 l 和 r 的值,计算和在 [l, r] 范围内的子集的数量并返回。

时间复杂度: O(n * s)

空间复杂度: O(r)

class Solution:
    def countSubMultisets(self, nums: List[int], l: int, r: int) -> int:
        s = sum(nums)  # 计算数组总和
        if s < l:  # 如果总和小于l,直接返回0
            return 0
        r, S = min(r, s), -1  # 调整r为不超过s的值,初始化S为-1
        c, m = Counter(nums), 10**9 + 7  # 计数器和模数
        k = s + 1 >> 1  # 中值
        if l >= k:  # 根据l和r的值调整范围
            l, r = s - r, s - l
        elif r >= k:
            l, r = sorted((l - 1, s - r - 1))
            S = reduce(lambda i, j: i * j % m, (v + 1 for v in c.values()))  # 计算所有组合的和
            if r < 0:
                return S
        s, r1 = 1, r + 1
        a = [0] * r1  # 动态规划数组初始化
        a[0] = c[0] + 1  # 基于0的元素的初始计数
        del c[0]  # 删除0元素
        for v, k in sorted((k * v, k) for k, v in c.items()):  # 遍历每个元素
            s = min(s + v, r1)
            a1 = a.copy()  # 复制当前数组状态
            v += k
            for i in range(k, s):  # 更新动态规划数组
                a[i] += a[i - k]
                if (j := i - v) >= 0:
                    a[i] -= a1[j]
                a[i] %= m
        if S < 0:
            return sum(a[l:]) % m  # 计算结果
        return (S - sum(a[:l + 1]) - sum(a)) % m  # 计算最终结果

Explore

在题解中,如果数组的总和s小于下界l,意味着无论如何选择数组中的元素,其组合的和都不可能达到l。因此,在总和s小于l的情况下,不存在任何子多重集的和能在区间[l, r]内,直接返回0是一种效率优化,避免了不必要的计算。

调整r为不超过s的值是为了确保考虑的子多重集的和不超过数组总和s的可能性。这样的调整避免了在动态规划中无意义的计算和内存使用,因为和大于s的子多重集本来就不存在。因此,这样的调整对算法的正确性无影响,但显著提高了算法的效率,特别是在r远大于s时。

在动态规划数组a中,`a[0] = c[0] + 1`表示包含0元素的子多重集的计数。`c[0]`是0元素在数组中的计数,即0元素出现的次数。`a[0] = c[0] + 1`意味着包括从选择0个0元素(即不包括任何0)到选择所有0元素的所有可能子多重集。因此,`a[0]`的初始值为选择0元素的所有可能方式的数量。

在算法中删除0元素的计数是为了简化接下来的动态规划处理。一旦初始化了包含0元素的情况(`a[0] = c[0] + 1`),后续的动态规划就不需要再考虑0元素的影响。这样做可以减少后续计算的复杂性,并防止0元素的重复处理。实际上,这是一种优化措施,使得算法更专注于处理非零元素的计数和组合。