[LeetCode] 2681. Power of Heroes

Problem

Image is linked to LeetCode problem page

Explanation

The power of a group of heroes is defined by the square of the maximum strength multiplied by the minimum strength: ``Power = (max)^2 \times min.`` To solve this efficiently, sort the array `[x_0, x_1, \dots , x_{n-1}]` in ascending order. Once sorted, for any subset with minimum index `i` and maximum index `j` `\{x_i , \dots , x_j\}` `(i < j)`, the power is simplified to ``Power = {x_j}^2 \cdot x_i``

To calculate the total power, analyze the contribution of each element `x_j` acting as the maximum. For any two heroes at indices `i` and `j`, there `2^{j-i-1}` unique subsets where `x_i` is the minimum and `x_j` is the maximum. 

Total power generated by subsets where `x_i` is the minimum and `x_j` is the maximum

If we fix `x_j` as the maximum, the total power it generates across all possible minimums `x_0` through `x_j` is: ``\begin{align} P_j &= 2^{j-0-1} (x_j^2 \cdot x_0) + 2^{j-1-1} (x_j^2 \cdot x_1) + \cdots + 2^{j-(j-1)-1} (x_j^2 \cdot x_{j-1}) + (x_j^2 \cdot x_j) \\ &= x_j^2 [(2^{j-1} \cdot x_0) + (2^{j-2} \cdot x_1) + \cdots + (2^0 \cdot x_{j-1}) + (1 \cdot x_j)] \end{align}``

The final term `(1 \cdot x_j)` represents the subset `\{x_j\}`, where it is both the minimum and the maximum.

For the same minimum `x_i`, maximum `x_{j+1}` has double the number of possible subsets compared to `x_{j}`

Comparing the sum of minimums for `x_j` to that of the next element `x_{j+1}`, the contribution of all minimums from `x_0` to `x_{j-1}` double as we increase the distance between the minimum and maximum heroes. The case where `x_{j}` is the minimum is an exception as the pair `\{x_j, x_{j+1}\}` only has one possible subset. We can capture this behavior using a running sum variable `S_j`, which maintains the weighted sum of all previous minimums such that `P_j = x_j^2 (S_j + x_j)`, and is updated each iteration as follows: ``S_{j+1} = 2S_{j} + x_{j}``

Remember to apply the modulo operator at each addition and multiplication step to prevent integer overflow.

Solution

#include <vector>
#include <algorithm>

class Solution {
public:
    int sumOfPower(vector<int>& nums) {
        const int MOD = 1e9 + 7;

        std::sort(nums.begin(), nums.end());

        long long power = 0;
        long long min_sum = 0;

        for(int x : nums) {
            long long max_sq = ((long long)x * x) % MOD;
            power = (power + (max_sq * (min_sum + x)) % MOD) % MOD;
            min_sum = (min_sum * 2 + x) % MOD;
        }
        return (int) power;
    }
};

Example

We are given ##nums=[1, 2, 3, 4, 5]##. The following table shows how the weighted minimum sum evolves as we fix each `x_j` as the maximum.

`x_j`	Subarrays where `x_j` is maximum	`P_j`	`S_j+x_j`
`1`	`[1]`	`1`	`0+1`
`2`	`[1,2]` `[2]`	`(2^2\cdot 1) + \\ (2^2 \cdot 2)`	`1+2`
`3`	`[1,3], [1,2,3]` `[2,3]` `[3]`	`2(3^2 \cdot1) + \\ 3^2 \cdot 2 + \\ 3^2 \cdot 3`	`(2 \cdot 1 + 2) +3`
`4`	`[1,4], [1,2,4], [1,3,4], [1,2,3,4]` `[2,4], [2,3,4]` `[3,4]` `[4]`	`4(4^2 \cdot 1) + \\ 2(4^2 \cdot 2) + \\ 4^2 \cdot 3 + \\ 4^2 \cdot 4`	`(4\cdot 1 + 2\cdot 2 + 3) +4`