[LeetCode]2281. Sum of Total Strength of Wizards 计算公式推导详解（JavaScript版)

题目描述

LeetCode原题链接：2281. Sum of Total Strength of Wizards

As the ruler of a kingdom, you have an army of wizards at your command.

You are given a 0-indexed integer array strength, where strength[i] denotes the strength of the ith wizard. For a contiguous group of wizards (i.e. the wizards' strengths form a subarray of strength), the total strength is defined as the product of the following two values:

• The strength of the weakest wizard in the group.
• The total of all the individual strengths of the wizards in the group.

Return the sum of the total strengths of all contiguous groups of wizards. Since the answer may be very large, return it modulo 109 + 7.

A subarray is a contiguous non-empty sequence of elements within an array.

Example 1:

Input: strength = [1,3,1,2]
Output: 44
Explanation: The following are all the contiguous groups of wizards:
- [1] from [1,3,1,2] has a total strength of min([1]) * sum([1]) = 1 * 1 = 1
- [3] from [1,3,1,2] has a total strength of min([3]) * sum([3]) = 3 * 3 = 9
- [1] from [1,3,1,2] has a total strength of min([1]) * sum([1]) = 1 * 1 = 1
- [2] from [1,3,1,2] has a total strength of min([2]) * sum([2]) = 2 * 2 = 4
- [1,3] from [1,3,1,2] has a total strength of min([1,3]) * sum([1,3]) = 1 * 4 = 4
- [3,1] from [1,3,1,2] has a total strength of min([3,1]) * sum([3,1]) = 1 * 4 = 4
- [1,2] from [1,3,1,2] has a total strength of min([1,2]) * sum([1,2]) = 1 * 3 = 3
- [1,3,1] from [1,3,1,2] has a total strength of min([1,3,1]) * sum([1,3,1]) = 1 * 5 = 5
- [3,1,2] from [1,3,1,2] has a total strength of min([3,1,2]) * sum([3,1,2]) = 1 * 6 = 6
- [1,3,1,2] from [1,3,1,2] has a total strength of min([1,3,1,2]) * sum([1,3,1,2]) = 1 * 7 = 7
The sum of all the total strengths is 1 + 9 + 1 + 4 + 4 + 4 + 3 + 5 + 6 + 7 = 44.

Example 2:

Input: strength = [5,4,6]
Output: 213
Explanation: The following are all the contiguous groups of wizards: 
- [5] from [5,4,6] has a total strength of min([5]) * sum([5]) = 5 * 5 = 25
- [4] from [5,4,6] has a total strength of min([4]) * sum([4]) = 4 * 4 = 16
- [6] from [5,4,6] has a total strength of min([6]) * sum([6]) = 6 * 6 = 36
- [5,4] from [5,4,6] has a total strength of min([5,4]) * sum([5,4]) = 4 * 9 = 36
- [4,6] from [5,4,6] has a total strength of min([4,6]) * sum([4,6]) = 4 * 10 = 40
- [5,4,6] from [5,4,6] has a total strength of min([5,4,6]) * sum([5,4,6]) = 4 * 15 = 60
The sum of all the total strengths is 25 + 16 + 36 + 36 + 40 + 60 = 213.

Constraints:

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

思路分析

题目的 hint1: “Consider the contribution of each wizard to the answer.” 给了我们很好的提示，那就是分别考虑每一个元素作为subArray的最小值时能够为结果贡献多少strength。不难想到，要使一个元素成为一个subArray中的最小值，我们需要分别向左向右去找边界，而边界位置就是向左向右第一个小于当前元素的位置（不存在的话就取原数组的左右边界）。比如，对于数组arr=[1, 4, 5, 2, 3, 4, 2, 1]，对于index=3位置的2来说，符合条件的的左右边界为index=0和index=7，即整个范围[1, 4, 5, 2, 3, 4, 2, 1]内的包含arr[3]=2的subArray。

通过观察，我们可以发现，所有subArray的和可以用左部分的subArray的和（[4, 5, 2]）与右部分的subArray的和（[3, 4, 2]）计算得到。即:

Sum of range(left, right) = leftSum * (right - mid) + rightSum * (mid - left)

接下来，让我们来看如何计算leftSum和rightSum。

先来看左部分，对于子数组[4, 5, 2]来说，leftSum中包含了1个4、2个5、3个2；好像有一定的规律？换一个子数组看看呢？如果是[1, 4, 5, 2]，那么它的leftSum就包含了1个1、2个4、3个5、4个2；如果是从index=0位置开始的子数组，那么leftSum就等于arr[0] * 1 + arr[1] * 2 + arr[2] * 3 + ... + arr[i] * (i + 1)，我们可以用一个prefixProductL数组来记录这个累加乘积。那如果不是从index=0位置开始的呢？再来看一个例子，对于子数组[2, 3, 4, 2]，其leftSum包含了1个2、2个3、3个4、4个2，可以发现，如果子数组的起始位置相较于index=0位置往后/右推移n位的话，包含每一位元素的个数就要递减n个，这里的每一位元素可以用前缀和prefixSumL来表示。

那么对于范围[start, end]的子数组来说:

leftSum = prefixProductL[end+1] - prefixProductL[start] - start * (prefixSumL[end+1] - prefixSumL[start])

同理，对于rightSum我们可以推导出:

rightSum = prefixProductR[start+1] - prefixProductR[end] - (arr.length - end) * (prefixSumR[start+1] - prefixSumR[end])

综上，对于任意位置i的元素来说，其对最终结果res的贡献为：

res += strength[i] * (leftSum * (right - mid) + rightSum * (mid - left))

这里找左右边界的时候可以使用单调递减栈来快速得到，如果当前遍历到的元素小于栈顶元素则入栈；否则当前元素作为右边界，先弹出一个栈顶元素作中间元素，此时新的栈顶元素就是左边界。

代码示例（JS)

var mod = 1000000007n;
var totalStrength = function(strength) {
    let len = strength.length;
    // record prefixSumL & prefixProductL
    let prefixSumL = new Array(len + 1), prefixProductL = new Array(len + 1);
    prefixSumL.fill(0n);
    prefixProductL.fill(0n);
    for(let i = 0; i < len; i++) {
        prefixSumL[i+1] = (prefixSumL[i] + BigInt(strength[i])) % mod;
        prefixProductL[i+1] = (prefixProductL[i] + BigInt((i + 1) * strength[i])) % mod;
    }
    // record prefixSumR & prefixProductR
    let prefixSumR = new Array(len + 1), prefixProductR = new Array(len + 1);
    prefixSumR.fill(0n);
    prefixProductR.fill(0n);
    for(let i = len - 1; i >= 0; i--) {
        prefixSumR[i] = (prefixSumR[i + 1] + BigInt(strength[i])) % mod;
        prefixProductR[i] = (prefixProductR[i + 1] + BigInt((len - i) * strength[i])) % mod;
    }
    // use monotonic stack to find boundaries of current value
    let stack = [], res = 0n;
    for(let right = 0; right <= len; right++) {
        while(stack.length && (right == len || strength[stack[stack.length-1]] >= strength[right])) {
            let mid = stack.pop(), left = stack.length ? (stack[stack.length-1] + 1) : 0;
            let leftSum = (mod + prefixProductL[mid+1] - prefixProductL[left] - BigInt(left) * (prefixSumL[mid+1] - prefixSumL[left])) % mod;
            let rightSum = (mod + prefixProductR[mid+1] - prefixProductR[right] - BigInt(len - right) * (prefixSumR[mid+1] - prefixSumR[right])) % mod;
            let total = (leftSum * BigInt(right - mid) + rightSum * BigInt(mid - left + 1)) % mod;
            // update final result
            res = (res + total * BigInt(strength[mid])) % mod;
        }
        stack.push(right);
    }
    return res;
};