Count of Sub-Multisets With Bounded Sum

Count of Sub-Multisets With Bounded Sum

You are given a 0-indexed array nums of non-negative integers, and two integers l and r.

Return the count of sub-multisets within nums where the sum of elements in each subset falls within the inclusive range of [l, r].

Since the answer may be large, return it modulo 109 + 7.

A sub-multiset is an unordered collection of elements of the array in which a given value x can occur 0, 1, ..., occ[x] times, where occ[x] is the number of occurrences of x in the array.

Note that:

• Two sub-multisets are the same if sorting both sub-multisets results in identical multisets.
• The sum of an empty multiset is 0.

 

Example 1:

Input: nums = [1,2,2,3], l = 6, r = 6
Output: 1
Explanation: The only subset of nums that has a sum of 6 is {1, 2, 3}.

Example 2:

Input: nums = [2,1,4,2,7], l = 1, r = 5
Output: 7
Explanation: The subsets of nums that have a sum within the range [1, 5] are {1}, {2}, {4}, {2, 2}, {1, 2}, {1, 4}, and {1, 2, 2}.

Example 3:

Input: nums = [1,2,1,3,5,2], l = 3, r = 5
Output: 9
Explanation: The subsets of nums that have a sum within the range [3, 5] are {3}, {5}, {1, 2}, {1, 3}, {2, 2}, {2, 3}, {1, 1, 2}, {1, 1, 3}, and {1, 2, 2}.

 

Constraints:

• 1 <= nums.length <= 2 * 104
• 0 <= nums[i] <= 2 * 104
• Sum of nums does not exceed 2 * 104.
• 0 <= l <= r <= 2 * 104

 

解题思路

  昨晚大概可以分析出是单调队列优化多重背包的做法，不过太久没做忘了。

  我们可以将原问题转换成求集合元素之和不超过 $r$ 的集合数量 $dp(r)$，以及集合元素之和不超过 $l-1$ 的集合数量 $dp(l-1)$，那么 $dp(r) - dp(l-1)$ 就是集合元素之和确切在 $[l,r]$ 范围内的集合数量了。

  为此我们就可以用动态规划进行求解 $dp(m)$，定义状态 $f(i,j)$ 表示从前 $i$ 个元素中选出总和不超过 $j$ 的所有方案的数量。根据是否选择第 $i$ 个数来划分状态，因此有状态转移方程 $$f(i,j) = f(i-1,j) + f(i-1,j-w_i)$$

  很明显时间复杂度是 $O(n \cdot m)$。但题目中还有个约束条件 $\sum\limits_{i=1}^{n}{w_i} \leq 2 \cdot 10^4$，假设原序列中不同的数最多有 $k$ 个，意味着有 $\dfrac{k(k+1)}{2} \leq 2 \cdot 10^4 \, \Rightarrow \, k < 200$。因此可以对原序列按照值来进行分组，那么最多会被分成 $200$ 组，问题就变成了分组背包的问题。如果直接用朴素的多重背包来求解的话，时间复杂度还是 $O(n \cdot m)$。假设 $k$ 是分组的数量，那么用单调队列优化后时间复杂度就优化到了 $O(k \cdot m)$。

  问题转换成多重背包后状态定义有所改变，$f(i,j)$ 表示从前 $i$ 组中选出总和不超过 $j$ 的所有方案的数量，根据选择第 $i$ 组的元素 $w_{i}$ 的数量进行状态划分，有状态转移方程 $$f(i,j) = \sum\limits_{k=0}^{s_{w_{i}}}{f(i-1,j-k \cdot w_i)}$$

  把式子展开，并发现之间的关联性：

$$\begin{align*} f(i,j) &= f(i-1,j) + f(i-1,j-w_i) + f(i-1,j-2 \cdot w_i) + \cdots + f(i-1,j-s_{w_i} \cdot w_i) \\\\ f(i,j-w_i) &= \qquad\qquad\quad\;\, f(i-1,j-w_i) + f(i-1,j-2 \cdot w_i) + \cdots + f(i-1,j-s_{w_i} \cdot w_i) + f(i-1,j-(s_{w_i}+1) \cdot w_i) \\\\ f(i,j-2 \cdot w_i) &= \qquad\qquad\quad\;\, f(i-1,j-w_i) + f(i-1,j-2 \cdot w_i) + \cdots + f(i-1,j-s_{w_i} \cdot w_i) + f(i-1,j-(s_{w_i}+1) \cdot w_i) + f(i-1,j-(s_{w_i}+2) \cdot w_i) \\ & \,\,\, \vdots \\ f\left(i, j - \left(\left\lfloor\frac{j}{w_i} \right\rfloor - 1 \right) \cdot w_i \right) &= f\left(i, j - \left(\left\lfloor\frac{j}{w_i} \right\rfloor - 1 \right) \cdot w_i\right) + f\left(i-1, j - \left\lfloor\frac{j}{w_i} \right\rfloor \cdot w_i\right) \\\\ f\left(i, j - \left\lfloor\frac{j}{w_i} \right\rfloor \cdot w_i\right) &= \qquad\qquad\qquad\qquad\qquad\qquad\quad\, f\left(i-1, j - \left\lfloor\frac{j}{w_i} \right\rfloor \cdot w_i\right) \end{align*}$$

  当固定 $w_i$ 后，对于所有可能的 $j \in [0,m]$，根据 $j$ 模 $w_i$ 的结果 $r = j \bmod w_i = j - \left\lfloor\frac{j}{w_i} \right\rfloor \cdot w_i$ 进行分类，那么就可以分成 $0 \sim w_i - 1$ 这么多类。对于固定的 $r$，从 $r$ 开始遍历，依次计算 $f(i,r), \, f(i, r + w_i), \, \ldots$，并在这个过程中维护一个大小不超过 $s_{w_i} + 1$ 的滑动窗口以及窗口内元素的和，对于那么对于 $f(i, r+k \cdot w_i)$ 其结果就是窗口内元素的和（即对应的状态转移方程）。

  还有个细节就是不能有 $w_i = 0$ 的情况，否则 $f(i,j)$ 就可以转移到 $f(i,j)$ 本身，就会形成环不能 dp 了。由于 $0$ 对总和没有影响，所以我们先跳过 $w_i = 0$ 的情况，最后计算答案时再把含 $0$ 的情况考虑上，因此 $dp(m)$ 最终的答案就是 $f(k,m) \times (s_0 +1)$。 

  AC 代码如下，时间复杂度为 $O(k \cdot m)$：

class Solution {
public:
    int mod = 1e9 + 7;
    
    int countSubMultisets(vector<int>& nums, int l, int r) {
        sort(nums.begin(), nums.end());
        vector<int> cnt(nums.back() + 1);
        for (auto &x : nums) {  // 统计每个值的个数
            cnt[x]++;
        }
        nums.erase(unique(nums.begin(), nums.end()), nums.end());   // 去重分组
        if (!nums[0]) nums.erase(nums.begin()); // 不考虑0的情况
        int n = nums.size();    // 分组的数量
        function<int(int)> dp = [&](int m) {
            if (m < 0) return 0;
            vector<vector<int>> f(n + 1, vector<int>(m + 1));
            for (int i = 0; i <= m; i++) {
                f[0][i] = 1;
            }
            for (int i = 1; i <= n; i++) {
                int w = nums[i - 1];
                for (int j = 0; j < w; j++) {
                    int s = 0;
                    for (int k = j; k <= m; k += w) {
                        if (k - (cnt[w] + 1) * w >= 0) s = (s - f[i - 1][k - (cnt[w] + 1) * w] + mod) % mod;    // 超出窗口大小
                        s = (s + f[i - 1][k]) % mod;
                        f[i][k] = s;
                    }
                }
            }
            int ret = (f[n][m] * (cnt[0] + 1ll)) % mod; // 再把0的情况考虑进来
            return ret;
        };
        return (dp(r) - dp(l - 1) + mod) % mod;
    }
};

 

参考资料

  第 115 场力扣夜喵双周赛：https://leetcode.cn/circle/discuss/maICM3/view/zXw3Ko/

  力扣第115场双周赛：https://www.bilibili.com/video/BV1H34y1g7k7/