2281. Sum of Total Strength of Wizards

Question

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 i<sup>th</sup> 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 10<sup>9</sup><span> </span>+ 7.

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

Solution

参考连接

根据题意，我们需要的结果是所有子数列乘以该子数列中最小值之和。

数组strength[]中的每一个元素，都充当一定范围内的子数列的最小值。因此我们需要**确定每个元素充当最小值的范围(left, right)**。然后计算这个元素可以组成的子数列的和与这个元素的乘积。

我们的问题相当于确定下一个更小元素，而单调栈专门处理这一类问题——“下一个更大/更小元素”
496. Next Greater Element I

通过单调栈，我们可以确定最小值为strength[i]的子数列的范围(left, right)，并将其保存在两个数组left[]与right[]中。

最后通过前缀和的计算，来确定(left, right)范围内的子数列之和。

单调栈

新建两个数组left[]， right[]保存以strength[i]为最小值的范围。

维护两个单调栈，保证下列关系：
strength[left] < strength[i]
strength[right] <= strength[i]

其中，left取小于号，而right可以等于是为了取值不重复。

注意，单调栈保存的是下标，这样使用起来更加灵活。

前缀和

为了得到范围（left, right）内的所有子数列之和。我们可以观察left, i, right之间的关系。

…left-1, left, left + 1, left + 2, … i-1, i, i+1, … right-1, right, right+1…

• 开始于 left+1的子数列:
  sum(left+1, … i) = prefix[i + 1] - prefix[left + 1]
  sum(left+1, … i+1) = prefix[i + 2] - prefix[left + 1]
  …
  sum(left+1, … right-1) = prefix[right] - prefix[left + 1]
• 开始于left+2的子数列:
  sum(left+2, … i) = prefix[i + 1] - prefix[left + 2]
  sum(left+2, … i+1) = prefix[i + 2] - prefix[left + 2]
  …
  sum(left+2, … right-1) = prefix[right] - prefix[left + 2]
  …
• 开始于i的子数列:
  sum(i, … i) = prefix[i + 1] - prefix[i]
  sum(i, … i+1) = prefix[i + 2] - prefix[i]
  …
  sum(i, … right-1) = prefix[right] - prefix[i]

合并这些式子，我们可以得到：

• 正数 部分:
  (prefix[i + 1] + prefix[i + 2] + ... + prefix[right]) * (i - left)
• 负数 部分:
  (prefix[left + 1] + prefix[left + 2] + ... + prefix[i]) * (right - i)

最后我们可以提前计算前缀和的前缀和来优化计算步骤，注意计算时需要将各个部分模除以保证不超过范围。
正数部分模除后需要加上MOD以保证其减去负数部分的结果大于0。

Code

class Solution {
    public int totalStrength(int[] strength) {
        long MOD = 1_000_000_007;
        int n = strength.length;
        
        Deque<Integer> mono = new ArrayDeque<>();
        int[] left = new int[n], right = new int[n];
        
        for(int i = 0; i < strength.length; i++){   //单调栈，记录比当前元素strength[i]小的第一个左侧位置strength[left]
            while(!mono.isEmpty() && strength[mono.peek()] >= strength[i]){
                mono.pop();
            }
            left[i] = mono.isEmpty() ? -1 : mono.peek();
            mono.push(i);
        }
        mono.clear();
        for(int i = n-1; i >=0; i--){   //单调栈，记录比当前元素strength[i]小的第一个右侧位置strength[right]
            while(!mono.isEmpty() && strength[mono.peek()] > strength[i]){
                mono.pop();
            }
            right[i] = mono.isEmpty() ? n : mono.peek();
            mono.push(i);
        }
        
        long[] prefix = new long[n+1], prefix_sum = new long[n+2];
        
        for(int i = 1; i <= strength.length; i++){  //计算前缀和
            prefix[i] = (prefix[i-1] + strength[i-1]) % MOD;
        }
        
        for(int i = 2; i <= strength.length+1; i++){    //计算前缀和的前缀和
            prefix_sum[i] = (prefix_sum[i-1] + prefix[i-1]) % MOD;
        }
        
        long res = 0;
        
        for(int i = 0; i < n; i++){
            res += ((prefix_sum[right[i] + 1] - prefix_sum[i + 1]) * (i - left[i]) % MOD + MOD -    //这里加MOD是为了保证前项大于后项，防止出现负数
                    (prefix_sum[i + 1] - prefix_sum[left[i] + 1]) * (right[i] - i) % MOD ) * strength[i];
            res %= MOD;
        }
        
        return (int) res;
    }
}