2281. 巫师的总力量和（单调栈+前缀和的前缀和）

2281. 巫师的总力量和

题目描述

示例

数据范围

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

思路

参考灵茶山艾府

• 以每个点作为最小值，计算包含该最小值的所有区间所造成的贡献。
• 设左右边界 $[L,R]$（即为包含该元素为最小值的左右区间端点），为了避免重复计算，考虑左侧严格小于当前元素的最近位置 $L-1$，以及右侧 小于等于当前元素的最近元素位置 $R+1$。
• 设当前枚举的巫师的能力值为 $v$，那么他对答案产生的贡献是 $v$ 乘上在左右边界 $[L,R]$内的所有包含 $v$ 的子数组的元素和的和。

---

如何计算子数组的元素和的和？

• 不妨固定右端点，移动左端点来考虑。

设子数组右端点为 $r$，左端点为 $l$，当前枚举的元素下标为 $i$，那么有 $l<=i<=r$

设 $strength$ 数组的前缀和为 $s$，其中 $s[i]=\sum _{j=0}^{i-1}strength[j]$，因此子数组 $[l,r]$的元素和可以表示为

$$s[r+1]-s[l]$$

在范围 $[L,R]$ 内的所有子数组的元素和的和可以表示为

$$\begin{array}{rl}& \sum _{r=i+1}^{R+1}\sum _{l=L}^{i}s[r]-s[l]\\ & =\sum _{r=i+1}^{R+1}((i-L+1)\cdot s[r]-\sum _{l=L}^{i}s[l])\\ & =(i-L+1)\cdot \sum _{r=i+1}^{R+1}s[r]-(R-i+1)\cdot \sum _{l=L}^{i}s[l]\end{array}$$

所以还需计算出前缀和 $s$ 的前缀和 $ss$， 其中 $ss[i]=\sum _{j=0}^{i-1}s[j]$，上式即为

$$(i-L+1)\cdot (ss[R+2]-ss[i+1])-(R-i+1)\cdot (ss[i+1]-ss[L])$$

在乘上 $v$ 即为当前巫师贡献，累加所有贡献即为答案。

---

注意的点

---

代码

class Solution {
public:
    const int mod = 1e9 + 7;
    #define ll long long

    int totalStrength(vector<int>& a) {
        
        int n = a.size();
        vector<int> L(n, -1), R(n, n);
        stack<int> s; 
        // 直接求出左右
        for (int i = 0; i < n; i++) {
            while(!s.empty() && a[s.top()] >= a[i]) R[s.top()] = i, s.pop();
            if(!s.empty()) L[i] = s.top();
            s.push(i); 
        }

        long long pre = 0;
        vector<int> ppre(n + 2);
        for (int i = 1; i <= n; i++) {
            pre += a[i - 1];
            ppre[i + 1] = (ppre[i] + pre) % mod;
        }

        int res = 0;
        for (int i = 0; i < n; i++) {
            ll l = L[i] + 1, r = R[i] - 1;
            ll ans = ((i - l + 1) * (ppre[r + 2] - ppre[i + 1]) - \
                      (r - i + 1) * (ppre[i + 1] - ppre[l])) % mod;
            res = (res + ans * a[i]) % mod;
        }
        return (res + mod) % mod; // 防止负数
    }
};