Number of Great Partitions

You are given an array nums consisting of positive integers and an integer k .

Partition the array into two ordered groups such that each element is in exactly one group. A partition is called great if the sum of elements of each group is greater than or equal to k .

Return the number of distinct great partitions. Since the answer may be too large, return it modulo ${10}^9 + 7$.

Two partitions are considered distinct if some element nums[i] is in different groups in the two partitions.

Example 1:

Input: nums = [1,2,3,4], k = 4
Output: 6
Explanation: The great partitions are: ([1,2,3], [4]), ([1,3], [2,4]), ([1,4], [2,3]), ([2,3], [1,4]), ([2,4], [1,3]) and ([4], [1,2,3]).

Example 2:

Input: nums = [3,3,3], k = 4
Output: 0
Explanation: There are no great partitions for this array.

Example 3:

Input: nums = [6,6], k = 2
Output: 2
Explanation: We can either put nums[0] in the first partition or in the second partition.
The great partitions will be ([6], [6]) and ([6], [6]).

Constraints:

• $1 \leq \text{nums.length}, k \leq 1000$
• $1 \leq \text{nums}[i] \leq {10}^9$

解题思路

　　首先还是容易想到用动态规划的，定义$f(i, j)$表示所有从前$i$个数中凑出总和为$j$的方案的个数，根据第$i$个物品选或不选进行状态划分，那么状态转移方程就是$f(i, j) = f(i-1, j) + f(i-1, j - \text{nums}[i])$。最后答案就是$\sum\limits_{i = k}^{s-k}{f(n, i)}$，其中$s$是所有数的总和。

　　可以发现$s$最大可以取到${10}^{12}$，因此这种做法的计算量可以达到${10}^{15}$级别。

　　现在倒过来想，能不能求出所有“不好分区”的数目，然后用总的分区数目减去“不好分区”的数目来得到答案呢？当然可以，如果求“不好分区”的数目那么状态的第二维就可以控制到$k$以内，时间复杂度就降到了$O(n \cdot k)$了。

　　为了方便下面用$f(i)$来表示$f(n, i)$，即从前$n$个物品中凑出总和为$i$的方案的个数。

　　“不好分区”分为$3$种情况：

1. 第$1$组的总和小于$k$，第$2$组的总和大于等于$k$。
2. 第$1$组的总和大于等于$k$，第$2$组的总和小于$k$。
3. 第$1$组和第$2$组的总和均小于$k$。

　　这里统一规定$f(i)$指的是第$1$组的总和，题目中这两个组明显是不可互换的，即 ([1,2,3], [4]) 和 ([4], [1,2,3]) 是两个不同的组。

　　其中对于第$1$、$2$种情况，$i$要同时满足$i < k$以及$s - i \geq k$，因此“不好分区”的个数都是$\sum\limits_{i = 0}^{\max\{ {k-1, s-k} \}} {f(i)}$，因此这两种的情况的总个数就是$\sum\limits_{i = 0}^{\max\{ {k-1, s-k} \}} {2 \cdot f(i)}$（$f(i)$一个是放第$1$组，一个是放第$2$组）。

　　对于第$3$种情况，$i$要同时满足$i < k$以及$s - i < k$，因此“不好分区”的个数是$\sum\limits_{i = \min\{ {0, s-k+1} \}}^{k-1} {f(i)}$。

　　最后因为每个物品可以选择放在第$1$组或第$2$组，有两种选择，因此总的分区数目是$2^n$。

　　最终答案就是$$2^n ~~- \sum\limits_{i = 0}^{\max\{ {k-1, s-k} \}} {f(i)} ~~- \sum\limits_{i = \min\{ {0, s-k+1} \}}^{k-1} {f(i)}$$

　　AC代码如下：

 1 class Solution {
 2 public:
 3     int mod = 1e9 + 7;
 4 
 5     int countPartitions(vector<int>& nums, int k) {
 6         int n = nums.size();
 7         vector<vector<int>> f(n + 1, vector<int>(k + 1));
 8         f[0][0] = 1;
 9         for (int i = 1; i <= n; i++) {
10             for (int j = 0; j <= k; j++) {  // 记得j从0开始，因为其中一组什么都不放也算是不合法的分区方案
11                 f[i][j] = f[i - 1][j];
12                 if (j >= nums[i - 1]) f[i][j] = (f[i][j] + f[i - 1][j - nums[i - 1]]) % mod;
13             }
14         }
15         int ret = 1;
16         long long s = 0;
17         for (int i = 0; i < n; i++) {
18             ret = ret * 2 % mod;
19             s += nums[i];
20         }
21         for (int i = 0; i < k; i++) {
22             long long t = s - i;    // 另外一组分区的总和
23             if (t >= k) ret = (ret - 2 * f[n][i] % mod) % mod;  // 第1、2种情况
24             else ret = (ret - f[n][i]) % mod;   // 第3种情况
25         }
26         return (ret + mod) % mod;
27     }
28 };

参考资料

　　计数 & 背包 dp：https://leetcode.cn/problems/number-of-great-partitions/solution/by-tsreaper-69ff/