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+7

class 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] = 1

            pt = 0
            for i in range(1, m):
                while pt < i and dist_from_center[i] - dist_from_center[pt] >= limit_lo:
                    pt += 1

                for v in range(k - 1):
                    f[i][v + 1] = (f[i][v + 1] + f_acc[pt][v]) % MOD
                for v in range(k):
                    f_acc[i + 1][v] = (f_acc[i][v] + f[i][v]) % MOD

            return f_acc[-1]  # 物品[0, m]之方法数

        ans = 0
        for 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]) % MOD
        return ans

原文地址：https://blog.csdn.net/weixin_43344005/article/details/140639424

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！