单调栈-复杂前缀和-2281. 巫师的总力量和

2022-05-28 22:13:32

问题描述：

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

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

问题求解：

首先我们需要去求得每个位置作为最小值的辐射范围，这一步可以通过单调栈获得，注意一边开一边闭来保证唯一性。

本题最大的难点在于如何高效的求和。

a x x x i x x x x b

   1 2 3  4 3 2 1

我们发现在计算的时候有阶梯的性质，因此可以考虑通过两个前缀和来计算对应的和。

具体来说:

对于左侧：sub1 = nums[a + 1] * 1 + nums[a + 2] * 2 ... + nums[i - 1] * (i - 1) = nums[a + 1] * ( a + 1) + nums[a + 2] * (a + 2) + ... + (nums[i - 1] * (i - 1) - a * (nums[a + 1] + nums[a + 2] ...)
对于右侧：sub2 = nums[i + 1] * (b - i - 1) + ... = b * (nums[i + 1] + nums[i + 2] + ...) - (nums[i + 1] * (i + 1) + ...)

因此我们可以通过建立两个presum来完成快速的求和运算。

class Solution:
    def totalStrength(self, nums: List[int]) -> int:
        n = len(nums)
        nums.insert(0, 0)
        presum1 = [0] * (n + 1)
        presum2 = [0] * (n + 1)
        for i in range(1, n + 1):
            presum1[i] = nums[i] + presum1[i - 1]
            presum2[i] = nums[i] * i + presum2[i - 1]
        # 左侧第一个更小的数字
        left = [0] * (n + 1)
        # 右侧第一个更小/等于的数字
        right = [n + 1] * (n + 1)
        stack = []
        for i in range(n, 0, -1):
            while stack and nums[stack[-1]] > nums[i]:
                stack.pop()
            right[i] = stack[-1] if stack else n + 1
            stack.append(i)
        stack = []
        for i in range(1, n + 1):
            while stack and nums[stack[-1]] >= nums[i]:
                stack.pop()
            left[i] = stack[-1] if stack else 0
            stack.append(i)
        res = 0
        mod = int(1e9 + 7)
        for i in range(1, n + 1):
            a = left[i]
            b = right[i]
            x = i - a
            y = b - i
            sub1 = (presum2[i - 1] - presum2[a] - a * (presum1[i - 1] - presum1[a]) % mod + mod) % mod
            sub1 = sub1 * y % mod
            sub2 = (b * (presum1[b - 1] - presum1[i]) % mod - (presum2[b - 1] - presum2[i]) + mod) % mod
            sub2 = sub2 * x % mod
            res = (res + (sub1 + sub2 + nums[i] * x * y % mod) * nums[i] % mod) % mod
        return res

标签：前缀,nums,min,2281,sum,巫师,值为,stack,mod
来源： https://www.cnblogs.com/hyserendipity/p/16322341.html