Question

2681. 英雄的力量

English Version

题目描述

给你一个下标从 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

解法

方法一：排序 + 数学

我们注意到，题目中涉及到子序列的最大值和最小值，数组中元素的顺序不影响最终的结果，因此我们可以先对数组进行排序。

接下来，我们枚举每个元素作为子序列的最小值，不妨记数组的每个元素为 $a_1, a_2, \cdots, a_n$。以 $a_i$ 作为最小值的子序列的贡献为：

$$ a_i \times (a_{i}^{2} + a_{i+1}^2 + 2 \times a_{i+2}^2 + 4 \times a_{i+3}^2 + \cdots + 2^{n-i-1} \times a_n^2) $$

我们注意到，每一个 $a_i$ 都会乘上 $a_i^2$，这一部分我们可以直接累加到答案中。剩下的部分，我们可以用一个变量 $p$ 来维护，初始时 $p = 0$。

接下来从右往左枚举 $a_i$，每次我们将 $a_i \times p$ 累加到答案中，然后令 $p = p \times 2 + a_i^2$。

枚举完所有的 $a_i$ 之后，返回答案即可。

时间复杂度 $O(n \times \log n)$，空间复杂度 $O(\log n)$。其中 $n$ 为数组的长度。

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

class Solution {
  public int sumOfPower(int[] nums) {
    final int mod = (int) 1e9 + 7;
    Arrays.sort(nums);
    long ans = 0, p = 0;
    for (int i = nums.length - 1; i >= 0; --i) {
      long x = nums[i];
      ans = (ans + (x * x % mod) * x) % mod;
      ans = (ans + x * p % mod) % mod;
      p = (p * 2 + x * x % mod) % mod;
    }
    return (int) ans;
  }
}

class Solution {
public:
  int sumOfPower(vector<int>& nums) {
    const int mod = 1e9 + 7;
    sort(nums.rbegin(), nums.rend());
    long long ans = 0, p = 0;
    for (long long x : nums) {
      ans = (ans + (x * x % mod) * x) % mod;
      ans = (ans + x * p % mod) % mod;
      p = (p * 2 + x * x % mod) % mod;
    }
    return ans;
  }
};

func sumOfPower(nums []int) (ans int) {
  const mod = 1e9 + 7
  sort.Ints(nums)
  p := 0
  for i := len(nums) - 1; i >= 0; i-- {
    x := nums[i]
    ans = (ans + (x*x%mod)*x) % mod
    ans = (ans + x*p%mod) % mod
    p = (p*2 + x*x%mod) % mod
  }
  return
}

function sumOfPower(nums: number[]): number {
  const mod = 10 ** 9 + 7;
  nums.sort((a, b) => a - b);
  let ans = 0;
  let p = 0;
  for (let i = nums.length - 1; i >= 0; --i) {
    const x = BigInt(nums[i]);
    ans = (ans + Number((x * x * x) % BigInt(mod))) % mod;
    ans = (ans + Number((x * BigInt(p)) % BigInt(mod))) % mod;
    p = Number((BigInt(p) * 2n + x * x) % BigInt(mod));
  }
  return ans;
}