好分区的数目

题目

2518. 好分区的数目

---

给你一个正整数数组 nums 和一个整数 k 。

分区 的定义是：将数组划分成两个有序的 组 ，并满足每个元素 恰好 存在于 某一个 组中。如果分区中每个组的元素和都大于等于 k ，则认为分区是一个好分区。

返回 不同 的好分区的数目。由于答案可能很大，请返回对 10^9 + 7 取余 后的结果。

如果在两个分区中，存在某个元素 nums[i] 被分在不同的组中，则认为这两个分区不同。

示例 1：

输入：nums = [1,2,3,4], k = 4
输出：6
解释：好分区的情况是 ([1,2,3], [4]), ([1,3], [2,4]), ([1,4], [2,3]), ([2,3], [1,4]), ([2,4], [1,3]) 和 ([4], [1,2,3]) 。

示例 2：

输入：nums = [3,3,3], k = 4
输出：0
解释：数组中不存在好分区。

示例 3：

输入：nums = [6,6], k = 2
输出：2
解释：可以将 nums[0] 放入第一个分区或第二个分区中。
好分区的情况是 ([6], [6]) 和 ([6], [6]) 。

提示：

• 1 <= nums.length, k <= 1000
• 1 <= nums[i] <= 10^9

题解

方法一：

思路

对于每个数nums[i]，我们有选与不选两种选择，对于选择的数可以放到第一个分区中，剩余的数放到第二个分区中，想到背包问题。
考虑到每个nums[i]的数值范围很大，我们可以考虑所有不合法的分法，然后让总的$2^n$种方案减去不合法的方案数就可以得到答案。

下面考虑不合法的方案，设分好后第一个分区的总和为a，第二个分区的总和为b。

不合法的分法有三类：

1. a < k 且 b < k
2. a < k 且 b >= k
3. a >= k 且 b < k

可以事先排除第1种情况，第1种情况没有可行的方案数。直接得到答案0

接下来可以用01背包计数

设dp[i][j]为前i个数中选取一些后，和为j的方案数。

初始化dp[0][0] = 1。

若不选第i个数，则有贡献dp[i-1][j]

若选第i个数，则有贡献dp[i-1][j-nums[i-1]]

因此状态转移为dp[i][j] = dp[i-1][j] + dp[i-1][j-nums[i-1]]

对于第二种不合法情况有$bad = \sum \limits_{i=0}^{k-1} dp[n][i]$种，而第三种情况和第二种情况具有对称性，因此也为bad种。

所以合法的方案数为$2^n-2*bad$种。

代码

class Solution {
public:
    using ll = long long;
    const ll M = 1e9+7;
    ll fpow(ll x, ll p) {
        ll rt = 1;
        while (p) {
            if (p&1) rt = rt*x%M;
            x = x*x%M;
            p>>=1;
        }
        return rt;
    }
    ll dp[1005][1005];
    int countPartitions(vector<int>& nums, int k) {
        if (accumulate(nums.begin(), nums.end(), 0LL) < 2*k) return 0;
        dp[0][0] = 1;
        int n = nums.size();
        for (int i=1; i<=n; i++) {
            for (int j=0; j<k; j++) {
                dp[i][j] = dp[i-1][j];
                if (j>=nums[i-1]) dp[i][j] = (dp[i][j]+dp[i-1][j-nums[i-1]])%M;
                // printf("%d,%d=%d\n",i,j,dp[i][j]);
            }
        }
        ll bad = 0;
        for (int i=0; i<k; i++) bad = (bad+dp[n][i])%M;
        return (fpow(2, n)-2*bad+M)%M;
    }
};