Skip to content

2518. Number of Great Partitions 👍

  • Time: $O(nk)$
  • Space: $O(k)$
 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
35
class Solution {
 public:
  int countPartitions(vector<int>& nums, int k) {
    const long sum = accumulate(nums.begin(), nums.end(), 0L);
    long ans = modPow(2, nums.size());
    vector<long> dp(k + 1);
    dp[0] = 1;

    for (const int num : nums)
      for (int i = k; i >= num; --i) {
        dp[i] += dp[i - num];
        dp[i] %= kMod;
      }

    // Substract the cases that're not satisfied.
    for (int i = 0; i < k; ++i)
      if (sum - i < k)  // Both group1 and group2 < k.
        ans -= dp[i];
      else
        ans -= dp[i] * 2;

    return (ans % kMod + kMod) % kMod;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x % kMod, (n - 1)) % kMod;
    return modPow(x * x % kMod, (n / 2)) % 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
class Solution {
  public int countPartitions(int[] nums, int k) {
    final int kMod = 1_000_000_007;
    final long sum = Arrays.stream(nums).asLongStream().sum();
    long ans = modPow(2, nums.length, kMod); // 2^n % kMod
    long[] dp = new long[k + 1];
    dp[0] = 1;

    for (final int num : nums)
      for (int i = k; i >= num; --i) {
        dp[i] += dp[i - num];
        dp[i] %= kMod;
      }

    // Substract the cases that're not satisfied.
    for (int i = 0; i < k; ++i)
      if (sum - i < k) // Both group1 and group2 < k.
        ans -= dp[i];
      else
        ans -= dp[i] * 2;

    return (int) ((ans % kMod + kMod) % kMod);
  }

  private int modPow(int x, int n, int kMod) {
    int res = 1;
    for (int i = 0; i < n; ++i)
      res = res * x % kMod;
    return res;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution:
  def countPartitions(self, nums: list[int], k: int) -> int:
    kMod = 1_000_000_007
    summ = sum(nums)
    ans = pow(2, len(nums), kMod)  # 2^n % kMod
    dp = [1] + [0] * k

    for num in nums:
      for i in range(k, num - 1, -1):
        dp[i] += dp[i - num]
        dp[i] %= kMod

    # Substract the cases that're not satisfied.
    for i in range(k):
      if summ - i < k:  # Both group1 and group2 < k.
        ans -= dp[i]
      else:
        ans -= dp[i] * 2

    return ans % kMod