【Leetcode】2281. Sum of Total Strength of Wizards

题目地址：

https://leetcode.com/problems/sum-of-total-strength-of-wizards/

给定给一个长$n$数组$A$，对于每个子数组$a$，求$f(a)=\min a\times \sum a$。求${\sum }_a(\min a\times \sum a)$。

设$A$的前缀和数组为$p$，对于$A[k]$，我们求出其右边第一个比其小的数$A[j]$（不存在则视为$j=n$）和左边第一个比其小于等于的数$A[i]$（不存在则视为$i=-1$），这两个数的位置可以用单调栈得出。那么$A[k]$可以作为$A[i:j]$的最小值。主要的问题是如何求$A[i+1:j-1]$的所有含$A[k]$的子数组的和的总和。设$p$的前缀和数组为$q$（$q[i]$是$p$的前$i$项的和，由于$p$的长度是$n+1$，$q$的长度则会变为$n+2$），则含$A[k]$的子数组的和的总和为：$$(k-i)(q[j+1]-q[k+1])-(j-k)(q[k+1]-q[i+1])$$可以这么理解，这个和即为以$A[k,k+1,...,j-1]$为结尾的前缀和减去$A[i,...,k-1]$结尾的前缀和的所有组合的总和，我们可以把被减数汇总，其会重复加$k-i$次，同理减数也可以汇总，其会重复加$j-k$次。代码如下：

class Solution {
 public:
  int totalStrength(vector<int>& v) {
    const int MOD = 1e9 + 7;
    int n = v.size();
    long sum = 0, ss[n + 2];
    memset(ss, 0, sizeof ss);
    for (int i = 1; i <= n; i++) {
      sum += v[i - 1];
      ss[i + 1] = (ss[i] + sum) % MOD;
    }
    int l[n], r[n];
    fill(l, l + n, 0);
    fill(r, r + n, n);
    stack<int> stk;
    for (int i = 0; i < n; i++) {
      while (stk.size() && v[stk.top()] > v[i]) {
        r[stk.top()] = i;
        stk.pop();
      }

      l[i] = stk.size() ? stk.top() : -1;
      stk.push(i);
    }

    int res = 0;
    for (int k = 0; k < n; k++) {
      sum = ((k - l[k]) * (ss[r[k] + 1] - ss[k + 1]) % MOD -
			(r[k] - k) * (ss[k + 1] - ss[l[k] + 1]) % MOD) % MOD;
      sum = (sum + MOD) % MOD;
      res = (res + sum * v[k]) % MOD;
    }

    return res;
  }
};

时空复杂度$O(n)$。