3082. 求出所有子序列的能量和

题目

给你一个长度为 n 的整数数组 nums 和一个正整数 k 。

一个整数数组的能量定义为和等于 k 的子序列的数目。

请你返回 nums 中所有子序列的 能量和 。

由于答案可能很大，请你将它对 109 + 7 取余后返回。

示例 1：

输入： nums = [1,2,3], k = 3

输出： 6

解释：

总共有 5 个能量不为 0 的子序列：

子序列 [1,2,3] 有 2 个和为 3 的子序列：[1,2,3] 和 [1,2,3] 。

子序列 [1,2,3] 有 1 个和为 3 的子序列：[1,2,3] 。

子序列 [1,2,3] 有 1 个和为 3 的子序列：[1,2,3] 。

子序列 [1,2,3] 有 1 个和为 3 的子序列：[1,2,3] 。

子序列 [1,2,3] 有 1 个和为 3 的子序列：[1,2,3] 。

所以答案为 2 + 1 + 1 + 1 + 1 = 6 。

示例 2：

输入： nums = [2,3,3], k = 5

输出： 4

解释：

总共有 3 个能量不为 0 的子序列：

子序列 [2,3,3] 有 2 个子序列和为 5 ：[2,3,3] 和 [2,3,3] 。

子序列 [2,3,3] 有 1 个子序列和为 5 ：[2,3,3] 。

子序列 [2,3,3] 有 1 个子序列和为 5 ：[2,3,3] 。

所以答案为 2 + 1 + 1 = 4 。

示例 3：

输入： nums = [1,2,3], k = 7

输出： 0

解释：不存在和为 7 的子序列，所以 nums 的能量和为 0 。

提示：

1 <= n <= 100

1 <= nums[i] <= 104

1 <= k <= 100

解法

题意比较绕，题目要求的是所有「子序列」的「子序列（满足条件）」个数。

具体来说, nums = [1, 2, 3, 4, 5, 6], k = 9, 当我们找到一个满足条件的子序列 [1, 3, 5], 那么它应该在所有包含它的序列中都被计算一次。

那么有多少序列会包含它呢，除了 1,3,5 必须包含外，还剩下 n-3 个元素，根据元素是否被包含，我们一共可以有 2^(n-3) 个序列。也就是说，当我们找到一个长度为 x 的满足条件的子序列时，一共有 2^(n-x) 个序列可以包含它。

问题转化为，我们需要找到各个长度的满足 sum == k 的子序列，具体其长度和数量来计算最终的答案。

我们使用背包 dp 来完成这一个过程，dp[i][j][p] 表示从下标为 [0,i] 的元素中选出 j 个元素构成序列，其中满足 sum == k 的序列的个数。根据定义, j 必须 <= i+1, 否则的话我们无法完成选择，对应值应当为 0.

转移方程：

dp[i][j][p] = dp[i-1][j][p] if i+1 >= j else 0
dp[i][j][p] += dp[i-1][j-1][p-nums[i]] if p-nums[i]>=0 else 0

代码


class Solution:
    def sumOfPower(self, nums: List[int], k: int) -> int:
        MOD = 10 ** 9 + 7
        n = len(nums)

        dp = [ [0] + [0] * k for j in range(n+1) ] 
        dp[0][0] = 1
        for i in range(n):
            for j in range(i+1, 0, -1):
                for p in range(k, -1, -1):
                    if i+1 < j:
                        dp[j][p] = 0
                    if p - nums[i] >= 0:
                        dp[j][p] += dp[j-1][p-nums[i]]
                        
        ans = 0
        for j in range(n+1):
            ans = (ans + dp[j][-1] * 2 ** (n-j)) % MOD
        return ans