2518-好分区的数目

Raphael Liu Lv10
  2023-09-13 00:47:26 2023-09-13 00:47:26 Created   2023-09-13 16:06:23 2023-09-13 16:06:23 Updated    

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

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

返回 不同 的好分区的数目。由于答案可能很大，请返回对 109 + 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] <= 109

视频讲解 已出炉，欢迎点赞三连，在评论区分享你对这场周赛的看法~

如果直接计算好分区的数目，我们可以用 01 背包来做，但是背包容量太大，会超时。

正难则反，我们可以反过来，计算坏分区的数目，即第一个组或第二个组的元素和小于 k 的方案数。根据对称性，我们只需要计算第一个组的元素和小于 k 的方案数，然后乘 2 即可。

注意，如果 nums 的所有元素之和小于 2k，则不存在好分区，我们可以特判这种情况，直接返回 0。如果不特判，计算坏分区会重复统计，导致错误的结果。

那么原问题就转换为「从 nums 中选择若干元素，使得元素和小于 k 的方案数」，这样用 01 背包就不会超时了。

具体来说，定义 f[i][j] 表示从前 i 个数中选择若干元素，和为 j 的方案数。

分类讨论：

  • 不选第 i 个数：f[i][j] = f[i-1][j]；
  • 选第 i 个数：f[i][j] = f[i-1][j-\textit{nums}[i]]。

因此 f[i][j] = f[i-1][j] + f[i-1][j-\textit{nums}[i]]。

初始值 f[0][0] = 1。

坏分区的数目 bad} =(f[n][0]+f[n][1]+\cdots+f[n][k-1])\cdot 2。

答案为所有分区的数目减去坏分区的数目，即 2^n-\textit{bad，这里 n 为 nums 的长度。

代码实现时，可以用倒序循环的技巧来压缩空间。

[sol1-Python3]
1
2
3
4
5
6
7
8
9
10
class Solution:
    def countPartitions(self, nums: List[int], k: int) -> int:
        if sum(nums) < k * 2: return 0
        MOD = 10 ** 9 + 7
        f = [0] * k
        f[0] = 1
        for x in nums:
            for j in range(k - 1, x - 1, -1):
                f[j] = (f[j] + f[j - x]) % MOD
        return (pow(2, len(nums), MOD) - sum(f) * 2) % MOD
[sol1-Java]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
    private static final int MOD = (int) 1e9 + 7;

    public int countPartitions(int[] nums, int k) {
        var sum = 0L;
        for (var x : nums) sum += x;
        if (sum < k * 2) return 0;
        var ans = 1;
        var f = new int[k];
        f[0] = 1;
        for (var x : nums) {
            ans = ans * 2 % MOD;
            for (var j = k - 1; j >= x; --j)
                f[j] = (f[j] + f[j - x]) % MOD;
        }
        for (var x : f)
            ans = (ans - x * 2 % MOD + MOD) % MOD; // 保证答案非负
        return ans;
    }
}
[sol1-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Solution {
    const int MOD = 1e9 + 7;
public:
    int countPartitions(vector<int> &nums, int k) {
        if (accumulate(nums.begin(), nums.end(), 0L) < k * 2) return 0;
        int ans = 1, f[k]; memset(f, 0, sizeof(f));
        f[0] = 1;
        for (int x : nums) {
            ans = ans * 2 % MOD;
            for (int j = k - 1; j >= x; --j)
                f[j] = (f[j] + f[j - x]) % MOD;
        }
        for (int x : f)
            ans = (ans - x * 2 % MOD + MOD) % MOD; // 保证答案非负
        return ans;
    }
};
[sol1-Go]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
func countPartitions(nums []int, k int) int {
	const mod int = 1e9 + 7
	sum := 0
	for _, x := range nums {
		sum += x
	}
	if sum < k*2 {
		return 0
	}
	ans := 1
	f := make([]int, k)
	f[0] = 1
	for _, x := range nums {
		ans = ans * 2 % mod
		for j := k - 1; j >= x; j-- {
			f[j] = (f[j] + f[j-x]) % mod
		}
	}
	for _, x := range f {
		ans -= x * 2
	}
	return (ans%mod + mod) % mod // 保证答案非负
}

复杂度分析

  • 时间复杂度：O(nk)，其中 n 为 nums 的长度。
  • 空间复杂度：O(k)。

相似题目

 Comments
On this page
2518-好分区的数目