golang刷leetcode：巫师的总力量和

作为国王的统治者，你有一支巫师军队听你指挥。

给你一个下标从 0 开始的整数数组 strength ，其中 strength[i] 表示第 i 位巫师的力量值。对于连续的一组巫师（也就是这些巫师的力量值是 strength 的 子数组），总力量 定义为以下两个值的 乘积 ：

巫师中 最弱 的能力值。

组中所有巫师的个人力量值 之和 。

请你返回 所有 巫师组的 总 力量之和。由于答案可能很大，请将答案对 109 + 7 取余 后返回。

子数组 是一个数组里 非空 连续子序列。

示例 1：

输入：strength = [1,3,1,2]

输出：44

解释：以下是所有连续巫师组：

- [1,3,1,2] 中 [1] ，总力量值为 min([1]) * sum([1]) = 1 * 1 = 1

- [1,3,1,2] 中 [3] ，总力量值为 min([3]) * sum([3]) = 3 * 3 = 9

- [1,3,1,2] 中 [1] ，总力量值为 min([1]) * sum([1]) = 1 * 1 = 1

- [1,3,1,2] 中 [2] ，总力量值为 min([2]) * sum([2]) = 2 * 2 = 4

- [1,3,1,2] 中 [1,3] ，总力量值为 min([1,3]) * sum([1,3]) = 1 * 4 = 4

- [1,3,1,2] 中 [3,1] ，总力量值为 min([3,1]) * sum([3,1]) = 1 * 4 = 4

- [1,3,1,2] 中 [1,2] ，总力量值为 min([1,2]) * sum([1,2]) = 1 * 3 = 3

- [1,3,1,2] 中 [1,3,1] ，总力量值为 min([1,3,1]) * sum([1,3,1]) = 1 * 5 = 5

- [1,3,1,2] 中 [3,1,2] ，总力量值为 min([3,1,2]) * sum([3,1,2]) = 1 * 6 = 6

- [1,3,1,2] 中 [1,3,1,2] ，总力量值为 min([1,3,1,2]) * sum([1,3,1,2]) = 1 * 7 = 7

所有力量值之和为 1 + 9 + 1 + 4 + 4 + 4 + 3 + 5 + 6 + 7 = 44 。

示例 2：

输入：strength = [5,4,6]

输出：213

解释：以下是所有连续巫师组：

- [5,4,6] 中 [5] ，总力量值为 min([5]) * sum([5]) = 5 * 5 = 25

- [5,4,6] 中 [4] ，总力量值为 min([4]) * sum([4]) = 4 * 4 = 16

- [5,4,6] 中 [6] ，总力量值为 min([6]) * sum([6]) = 6 * 6 = 36

- [5,4,6] 中 [5,4] ，总力量值为 min([5,4]) * sum([5,4]) = 4 * 9 = 36

- [5,4,6] 中 [4,6] ，总力量值为 min([4,6]) * sum([4,6]) = 4 * 10 = 40

- [5,4,6] 中 [5,4,6] ，总力量值为 min([5,4,6]) * sum([5,4,6]) = 4 * 15 = 60

所有力量值之和为 25 + 16 + 36 + 36 + 40 + 60 = 213 。

提示：

1 <= strength.length <= 105

1 <= strength[i] <= 109

解题思路：

暴力解：求范围内的最小值和范围内和，然后相加，复杂度O(n^ 2)，显然不符合要求

单调栈+前缀和：

1，首先假设位置i时当前考虑范围[L,R]内值最小的，我们如何求L，R位置呢？

2，答案是单调栈：由于有重复的元素，为了避免重复计算，我们左侧寻找第一个严格小于当前元素的位置，右侧选择第一个小于等于当前元素的位置。

3，以右侧为例，我们从数组最右侧开始，判断栈顶元素，如果栈顶元素比当前元素大，元素出栈，直到比当前元素小为止，或者空，记录栈顶元素对应位置，将当前元素入栈。

4，左侧类似，这样我们通过这两个辅助数组，可以很好定位，左右边界[L,R]

5,如果我们要求L，R范围内的和，我们可以用S[1-R]-S[1-L],即前缀和

6，本题需要计算范围内和的和，所以需要前缀和的和SS

7，枚举[L,R]内所有集合，计算每个集合元素的和的和，每个元素计算了多少遍呢？R-L+1

8，例如区间内元素是[7,8,9,6]

9,包含8的子区间有[8],[7,8],[8,9],[7,8,9],[8,9,6],[7,8,9,6],其中[7,8],[7,8,9],[7,8,9,6]在计算7的时候已经计算过，所以相应的，每个元素被计算的次数为R-l+1,其中l变化范围为L到R

10，转化为前缀和为(R-l+1)*S[l] (L<=l<R)

代码实现：

func totalStrength(strength []int) (ans int) {
  const mod int = 1e9 + 7

  n := len(strength)
  left := make([]int, n) // left[i] 为左侧严格小于 strength[i] 的最近元素位置（不存在时为 -1）
  st := []int{}
  for i, v := range strength {
    for len(st) > 0 && strength[st[len(st)-1]] >= v {
      st = st[:len(st)-1]
    }
    if len(st) > 0 {
      left[i] = st[len(st)-1]
    } else {
      left[i] = -1
    }
    st = append(st, i)
  }

  right := make([]int, n) // right[i] 为右侧小于等于 strength[i] 的最近元素位置（不存在时为 n）
  st = []int{}
  for i := n - 1; i >= 0; i-- {
    v := strength[i]
    for len(st) > 0 && strength[st[len(st)-1]] > v {
      st = st[:len(st)-1]
    }
    if len(st) > 0 {
      right[i] = st[len(st)-1]
    } else {
      right[i] = n
    }
    st = append(st, i)
  }

  s := make([]int, n+1) // 前缀和
  for i, v := range strength {
    s[i+1] = (s[i] + v) % mod
  }
  ss := make([]int, n+2) // 前缀和的前缀和
  for i, v := range s {
    ss[i+1] = (ss[i] + v) % mod
  }
  for i, v := range strength {
    l, r := left[i]+1, right[i]-1 // [l,r] 左闭右闭
    res := ((i-l+1)*(ss[r+2]-ss[i+1]) - (r-i+1)*(ss[i+1]-ss[l])) % mod
    ans = (ans + res*v) % mod // 累加贡献
  }
  return (ans + mod) % mod // 防止算出负数（因为上面 res 有个减法）
}

复制