LC 2518. 好分区的数目

https://leetcode.cn/contest/weekly-contest-325/problems/number-of-great-partitions/

这题稍微有点绕。我们可以计算出一个分区上 `dp[i][k]` (表示考虑到第i个元素，所有元素加起来和为k)的数量，但是我们却很难计算出 `分区所有元素和>=k` 的情况，因为这个动态规划的数量太大了。

换个思路，可以使用另外一个角度来进行计算：

  1. (G1, G2)所有可能的组合数量 `(2 ** n)`
  2. `sum(G1) < k` 的所有可能数量，假设是 `X`
  3. `sum(G2) < k` 的所有可能数量一样是 `X`
  4. `sum(G1) < k & sum(G2) < k` 的所有可能数量是 `Y` 的话
  5. 那么 `sum(G1) >=k & sum(G2)>=k` 的数量就是 `(2**n) - X - X + Y`

其中X, Y都可以使用一个动态规划计算出来。下面代码里面可能MOD没有怎么处理好，但是可以通过测试用例。 

class Solution:
    def countPartitions(self, nums: List[int], k: int) -> int:
        n = len(nums)

        def precompute():
            dp = [[0] * (k) for _ in range(n + 1)]
            dp[0][0] = 1
            for i in range(n):
                for j in range(0, k):
                    j2 = nums[i] + j
                    if j2 < k:
                        dp[i + 1][j2] += dp[i][j]
                    dp[i + 1][j] += dp[i][j]
            return dp

        def pow(a, b, MOD):
            ans = 1
            while b:
                if b & 0x1:
                    ans = ans * a
                    ans = ans % MOD
                a = (a * a) % MOD
                b = b >> 1
            return ans

        # (G1, G2) = (2**n) % MOD
        # G1<k = dp[n][j] j in [0, k)
        # G2<k = dp[n][j] j in [0, k)
        # G1<k & G2<k  = dp[n][j] j in [0,k) and (SUM-j) in [0,k)
        # (G1,G2) - (G1<k) - (G2<k) + (G1<k & G2<k) = (G1>=k & G2>=k)
        dp = precompute()
        SUM = sum(nums)
        MOD = 10 ** 9 + 7
        ans = pow(2, n, MOD)
        for j in range(0, k):
            ans -= 2 * dp[n][j]
        for j in range(0, k):
            if (SUM - j) < k:
                ans += dp[n][j]
        return ans % MOD