Skip to content

2902. Count of Sub-Multisets With Bounded Sum 👍

  • Time: $O(r\cdot\sqrt{\Sigma(\texttt{nums[i]})})$
  • Space: $O(r)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
 public:
  int countSubMultisets(vector<int>& nums, int l, int r) {
    constexpr int kMod = 1'000'000'007;
    // dp[i] := the number of submultisets of `nums` with sum i
    vector<long> dp(r + 1);
    dp[0] = 1;
    unordered_map<int, int> count;

    for (const int num : nums)
      ++count[num];

    const int zeros = count[0];
    count.erase(0);

    for (const auto& [num, freq] : count) {
      // stride[i] := dp[i] + dp[i - num] + dp[i - 2 * num] + ...
      vector<long> stride = dp;
      for (int i = num; i <= r; ++i)
        stride[i] += stride[i - num];
      for (int i = r; i > 0; --i)
        if (i >= num * (freq + 1))
          // dp[i] + dp[i - num] + dp[i - freq * num]
          dp[i] = (stride[i] - stride[i - num * (freq + 1)]) % kMod;
        else
          dp[i] = stride[i] % kMod;
    }

    long ans = 0;
    for (int i = l; i <= r; ++i)
      ans = (ans + dp[i]) % kMod;
    return ((zeros + 1) * ans) % kMod;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
  public int countSubMultisets(List<Integer> nums, int l, int r) {
    final int kMod = 1_000_000_007;
    // dp[i] := the number of submultisets of `nums` with sum i
    long[] dp = new long[r + 1];
    dp[0] = 1;
    Map<Integer, Integer> count = new HashMap<>();

    for (final int num : nums)
      count.merge(num, 1, Integer::sum);

    final int zeros = count.containsKey(0) ? count.remove(0) : 0;

    for (Map.Entry<Integer, Integer> entry : count.entrySet()) {
      final int num = entry.getKey();
      final int freq = entry.getValue();
      // stride[i] := dp[i] + dp[i - num] + dp[i - 2 * num] + ...
      long[] stride = dp.clone();
      for (int i = num; i <= r; ++i)
        stride[i] += stride[i - num];
      for (int i = r; i > 0; --i)
        if (i >= num * (freq + 1))
          // dp[i] + dp[i - num] + dp[i - freq * num]
          dp[i] = (stride[i] - stride[i - num * (freq + 1)]) % kMod;
        else
          dp[i] = stride[i] % kMod;
    }

    long ans = 0;
    for (int i = l; i <= r; ++i)
      ans = (ans + dp[i]) % kMod;
    return (int) (((zeros + 1) * ans) % kMod);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution:
  def countSubMultisets(self, nums: list[int], l: int, r: int) -> int:
    kMod = 1_000_000_007
    # dp[i] := the number of submultisets of `nums` with sum i
    dp = [1] + [0] * r
    count = collections.Counter(nums)
    zeros = count.pop(0, 0)

    for num, freq in count.items():
      # stride[i] := dp[i] + dp[i - num] + dp[i - 2 * num] + ...
      stride = dp.copy()
      for i in range(num, r + 1):
        stride[i] += stride[i - num]
      for i in range(r, 0, -1):
        if i >= num * (freq + 1):
          # dp[i] + dp[i - num] + dp[i - freq * num]
          dp[i] = stride[i] - stride[i - num * (freq + 1)]
        else:
          dp[i] = stride[i]

    return (zeros + 1) * sum(dp[l:r + 1]) % kMod