LeetCode-day23-3098. 求出所有子序列的能量和

题目描述

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

一个 子序列的 能量 定义为子序列中 任意 两个元素的差值绝对值的 最小值 。

请你返回 nums 中长度 等于 k 的 所有 子序列的 能量和 。

由于答案可能会很大，将答案对 10⁹ + 7 取余 后返回。

示例

示例1：

输入：nums = [1,2,3,4], k = 3
输出：4
解释：
nums 中总共有 4 个长度为 3 的子序列：[1,2,3] ，[1,3,4] ，[1,2,4] 和 [2,3,4] 。能量和为 |2 - 3| + |3 - 4| + |2 - 1| + |3 - 4| = 4 。

示例2：

输入：nums = [2,2], k = 2
输出：0
解释：
nums 中唯一一个长度为 2 的子序列是 [2,2] 。能量和为 |2 - 2| = 0 。

示例3：

输入：nums = [4,3,-1], k = 2
输出：10
解释：
nums 总共有 3 个长度为 2 的子序列：[4,3] ，[4,-1] 和 [3,-1] 。能量和为 |4 - 3| + |4 - (-1)| + |3 - (-1)| = 10 。

思路

子序列问题的拆解 + 前缀和优化

代码

MOD = 10**9+7class Solution:def sumOfPowers(self, a: List[int], k: int) -> int:n = len(a)a.sort()def calc(dist_from_center: List[int], limit_lo: int) -> int:m = len(dist_from_center)  # 从中心点算起的距离f = [[0] * k for _ in range(m)]  # f[i][j]: 取到第i个元素时，拿j个物品的方法数f[0][1] = 1  # 背包问题方案数f_acc = [[0] * k for _ in range(m + 1)]  # f_acc[i][j]: 物品[0, i-1], 拿j物品的方法数f_acc[1][1] = 1pt = 0for i in range(1, m):while pt < i and dist_from_center[i] - dist_from_center[pt] >= limit_lo:pt += 1for v in range(k - 1):f[i][v + 1] = (f[i][v + 1] + f_acc[pt][v]) % MODfor v in range(k):f_acc[i + 1][v] = (f_acc[i][v] + f[i][v]) % MODreturn f_acc[-1]  # 物品[0, m]之方法数ans = 0for i in range(n):for j in range(i):min_diff = a[i] - a[j]  # 最小差值dist_left = [a[j] - a[k] for k in range(j, -1, -1)]  # 注意取距离中心点的距离，要包含自己!f_left = calc(dist_left, min_diff + 1)  # 左右随便找一个不包含，都不会重复dist_right = [a[k] - a[i] for k in range(i, n)]f_right = calc(dist_right, min_diff)for x in range(1, k):  # 枚举左右取多少，左右至少取一个ans = (ans + min_diff * f_left[x] * f_right[k - x]) % MODreturn ans