巫师的总力量和

标签: 数组 前缀和 单调栈

难度: Hard

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

给你一个下标从 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

Submission

运行时间: 334 ms

内存: 34.2 MB

class Solution:
    def totalStrength(self, strength: List[int]) -> int:
        #0x3
        n = len(strength)
        left,right,st = [-1]*n,[n]*n,[]
        for i, v in enumerate(strength):
            while st and strength[st[-1]] >= v:
                right[st.pop()] =i
            if st:
                left[i] = st[-1]
            st.append(i)
        
        ss = list(accumulate(accumulate(strength,initial =0),initial = 0))

        ans = 0
        for i ,v in enumerate(strength):
            l,r = left[i]+1,right[i]-1
            tot = (i-l+1)*(ss[r+2]-ss[i+1])-(r-i+1)*(ss[i+1]-ss[l])
            ans += v* tot
        return ans%(1000000007)

Explain

此题解使用了单调栈来寻找每个元素在数组中作为最小值能扩展到的最远左右边界。首先,通过遍历数组使用单调栈计算每个元素的左右边界,其中左边界表示当前元素左侧第一个比它小的元素的位置,右边界表示当前元素右侧第一个比它小的元素的位置。此外,使用前缀和技术来快速计算任意子数组的和,避免重复计算带来的性能损失。最后,结合前缀和和左右边界,快速计算出每个元素作为最小值时,所有包含它的子数组的总力量值的和,并累加起来得到最终结果。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def totalStrength(self, strength: List[int]) -> int:
        n = len(strength)
        left, right, st = [-1] * n, [n] * n, []
        for i, v in enumerate(strength):
            while st and strength[st[-1]] >= v:
                right[st.pop()] = i
            if st:
                left[i] = st[-1]
            st.append(i)
        
        ss = list(accumulate(accumulate(strength, initial=0), initial=0))
        
        ans = 0
        for i, v in enumerate(strength):
            l, r = left[i] + 1, right[i] - 1
            tot = (i - l + 1) * (ss[r + 2] - ss[i + 1]) - (r - i + 1) * (ss[i + 1] - ss[l])
            ans += v * tot
        return ans % (1000000007)

Explore

单调栈是一种可以维护数组元素顺序和值的单调性的数据结构。在这道题中,使用一个递增的单调栈可以帮助我们有效地找到每个元素左侧和右侧第一个比它小的元素的位置。当遍历到新的元素时,栈从顶部开始检查,如果栈顶元素大于等于当前元素,栈顶元素出栈并更新其右边界为当前元素的索引。这样可以保证每个元素被处理时,栈中所有元素都是单调递增的,从而快速找到边界。

在计算左右边界的过程中,停止并设置边界的时候是因为我们需要找到每个元素作为最小值可以扩展到的最远范围。一旦在某一方向找到了一个更小的元素,这个元素就定义了当前元素作为最小值的边界,因为在这之外的任何较大范围中,当前元素就不再是最小的了。这样的处理保证了正确地计算每个元素作为最小值时的子数组范围。

前缀和数组ss是通过两次使用accumulate函数构建的。第一次accumulate计算了数组strength的前缀和,即到当前元素为止的所有元素的累加和。第二次accumulate则是基于第一次的结果,再次计算累加和,这样得到的是数组中任意子数组和的累加值。这种方法可以在常数时间内计算出任意子数组的和,大大提高了效率。

解法中的数学公式用于计算每个元素作为最小值时的子数组总力量和。公式是:`tot = (i - l + 1) * (ss[r + 2] - ss[i + 1]) - (r - i + 1) * (ss[i + 1] - ss[l])`。这里,`(i - l + 1)` 和 `(r - i + 1)` 分别表示从当前元素i到其左边界l和右边界r的元素数量。`ss[r + 2] - ss[i + 1]` 计算的是从i到r的子数组的和,而 `ss[i + 1] - ss[l]` 计算的是从l到i的子数组的和。整体公式考虑了所有包含当前元素并且当前元素是最小值的子数组,并计算出其总力量和。