LeetCode 2681. Power of Heroes Solution in Java, C++, Python & More | Explanation + Code

Description

You are given a 0-indexed integer array nums representing the strength of some heroes. The power of a group of heroes is defined as follows:

• Let i0, i1, ... ,ik be the indices of the heroes in a group. Then, the power of this group is max(nums[i0], nums[i1], ... ,nums[ik])2 * min(nums[i0], nums[i1], ... ,nums[ik]).

Return the sum of the power of all non-empty groups of heroes possible. Since the sum could be very large, return it modulo 109 + 7.

 

Example 1:

Input: nums = [2,1,4]
Output: 141
Explanation: 
1st group: [2] has power = 22 * 2 = 8.
2nd group: [1] has power = 12 * 1 = 1. 
3rd group: [4] has power = 42 * 4 = 64. 
4th group: [2,1] has power = 22 * 1 = 4. 
5th group: [2,4] has power = 42 * 2 = 32. 
6th group: [1,4] has power = 42 * 1 = 16. 
​​​​​​​7th group: [2,1,4] has power = 42​​​​​​​ * 1 = 16. 
The sum of powers of all groups is 8 + 1 + 64 + 4 + 32 + 16 + 16 = 141.

Example 2:

Input: nums = [1,1,1]
Output: 7
Explanation: A total of 7 groups are possible, and the power of each group will be 1. Therefore, the sum of the powers of all groups is 7.

 

Constraints:

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

Solutions

Solution 1: Sorting + Mathematics

We notice that the problem involves the maximum and minimum values of a subsequence, and the order of elements in the array does not affect the final result. Therefore, we can sort the array first.

Next, we enumerate each element as the minimum value of the subsequence. Let's denote each element of the array as a₁, a₂, …, a_n. The contribution of the subsequence with a_i as the minimum value is:

a_i × (a_i^{2} + a_i+1^2 + 2 × a_i+2^2 + 4 × a_i+3^2 + … + 2^n-i-1 × a_n^2)

We notice that each a_i will be multiplied by a_i^2, which we can directly add to the answer. For the remaining part, we can maintain it with a variable p, initially set to 0.

Then, we enumerate a_i from right to left. Each time, we add a_i × p to the answer, and then set p = p × 2 + a_i^2.

After enumerating all a_i, return the answer.

The time complexity is O(n × log n), and the space complexity is O(log n). Here, n is the length of the array.

PythonJavaC++GoTypeScript

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(code-box)

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; } }(code-box)

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; } };(code-box)

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 }(code-box)

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; }(code-box)