英雄的力量

标签: 数组 数学 前缀和 排序

难度: Hard

给你一个下标从 0 开始的整数数组 nums ，它表示英雄的能力值。如果我们选出一部分英雄，这组英雄的 力量 定义为：

  • i0 ，i1 ，... ik 表示这组英雄在数组中的下标。那么这组英雄的力量为 max(nums[i0],nums[i1] ... nums[ik])2 * min(nums[i0],nums[i1] ... nums[ik]) 。

请你返回所有可能的 非空 英雄组的 力量 之和。由于答案可能非常大，请你将结果对 109 + 7 取余。

示例 1：

输入：nums = [2,1,4]
输出：141
解释：
第 1 组：[2] 的力量为 22 * 2 = 8 。
第 2 组：[1] 的力量为 12 * 1 = 1 。
第 3 组：[4] 的力量为 42 * 4 = 64 。
第 4 组：[2,1] 的力量为 22 * 1 = 4 。
第 5 组：[2,4] 的力量为 42 * 2 = 32 。
第 6 组：[1,4] 的力量为 42 * 1 = 16 。
第​ ​​​​​​7 组：[2,1,4] 的力量为 42​​​​​​​ * 1 = 16 。
所有英雄组的力量之和为 8 + 1 + 64 + 4 + 32 + 16 + 16 = 141 。

示例 2：

输入：nums = [1,1,1]
输出：7
解释：总共有 7 个英雄组，每一组的力量都是 1 。所以所有英雄组的力量之和为 7 。

提示：

  • 1 <= nums.length <= 105
  • 1 <= nums[i] <= 109

Submission

运行时间: 100 ms

内存: 26.2 MB

class Solution:
    def sumOfPower(self, nums: List[int]) -> int:
        MOD = 10 ** 9 + 7
        nums.sort()
        ans = s = 0
        for x in nums:
            ans = (ans + x*x*(x+s)) % MOD
            s = (s*2 + x) % MOD
        return ans

Explain

本题解首先对数组进行排序，确保数组元素按升序排列。排序之后，数组中的每个元素x在某个位置i时，它在前i个元素中是最大的。因此，对于数组中的每个元素x，我们可以计算以x作为最大值的所有子集的力量和，因为x左侧的所有元素都不超过x。使用变量s来累加之前所有元素的和（考虑了每个元素出现的次数），这样可以快速计算出包含x的所有子集的最小值之和。对于每个元素x，其贡献的力量可以表示为x^2 * (x + s)，其中s是x左侧所有元素的和，每次迭代时更新s为s*2 + x。这种方法避免了直接枚举所有子集，通过数学推导简化了计算过程。

时间复杂度: O(n log n)

空间复杂度: O(1)

class Solution:
    def sumOfPower(self, nums: List[int]) -> int:
        MOD = 10 ** 9 + 7  # 定义取余操作的模
        nums.sort()  # 对输入数组排序
        ans = s = 0  # 初始化结果和前缀和
        for x in nums:  # 遍历排序后的数组
            ans = (ans + x*x*(x+s)) % MOD  # 计算包含当前元素x的所有子集的力量和
            s = (s*2 + x) % MOD  # 更新前缀和，考虑每个元素在新的子集中出现两次
        return ans  # 返回最终结果

Explore

排序的主要目的是为了确保当遍历到数组中的某个元素x时，它是当前遍历过的所有元素中的最大值。这样，可以保证在计算子集的最小值之和时，如果子集包含x，那么x就是该子集的最小值。这种方法避免了复杂的比较操作，并使得计算每个元素x作为子集最小值的情况变得直接且清晰。

在该算法中，变量`s`表示在当前元素x之前的所有元素的和。随着数组的遍历，每一个新元素都会存在于新的子集中，而之前的每个元素会在新的子集中出现两倍（即原有的子集加上新元素形成的新子集）。因此，更新过程`s*2 + x`意味着将之前所有元素的总和翻倍（考虑到新的子集的形成），然后加上当前元素x本身，以此来更新s的值，反映出所有包含当前元素x及其之前元素的子集的元素总和。

表达式`x*x*(x+s)`是基于以下数学原理推导出来的：对于数组中的某个元素x，在x之前的任何元素都不会超过x，因此如果包含x的子集选择x作为最小值，那么这个子集的力量就是x乘以该子集的元素总和。这里，x+s表示包含x的所有可能子集的元素总和（s是之前元素的总和，x是当前元素）。因此，x乘以(x+s)就得到了包含x的所有子集的力量总和。再乘以x是因为每个子集的力量是它的最小元素乘以子集的总和。

通过排序，我们确保了每次处理元素x时，它是到目前为止遇到的最大元素。由于排序保证了x之前的任何元素都不大于x，因此在考虑包含x的所有可能子集时，x将是这些子集中的最小值。这样，我们就可以简单地通过累积之前元素的总和并用这些总和来计算包含当前元素的所有子集的力量值，从而确保计算的准确性和高效性。