2281-巫师的总力量和

Raphael Liu Lv10
  2023-09-13 00:47:26 2023-09-13 00:47:26 Created   2023-09-13 16:06:25 2023-09-13 16:06:25 Updated    

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

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

本题 视频讲解 已出炉，额外介绍了单调栈的原理，欢迎点赞三连~

提示 1-1

枚举每位巫师，假设他是最弱的巫师，那么他能在哪些子数组中？

提示 1-2

左右边界最远能到哪？具体地，这些子数组的左边界的最小值是多少，右边界的最大值是多少？

提示 1-3

用单调栈来计算左右边界。

不了解单调栈的同学可以看一下 496. 下一个更大元素 I 。对于本题，我们需要求的是下一个更小元素。

提示 1-4

注意本题是可能有重复元素的，这会对最终答案的计算产生什么影响？

提示 1-5

设左右边界为 [L,R]。

为了避免重复计算，我们可以考虑左侧严格小于当前元素的最近元素位置 L-1，以及右侧小于等于当前元素的最近元素位置 R+1。

以示例 1 中的数组 [1,3,1,2] 为例，如果左右两侧都是找严格小于，那么第一个 1 和第二个 1 算出来的边界范围都是一样的（都是整个数组），这就重复统计了，为了避免这种情况，可以把某一侧改为小于等于，比如把右侧改成小于等于，那么第一个 1 算出来的右边界不会触及或越过第二个 1，这样就能避免重复统计同一个子数组。

提示 2-1

设当前枚举的巫师的能力值为 v，那么他对答案产生的贡献是 v 乘上在左右边界 [L,R] 内的所有包含 v 的子数组的元素和的和。

提示 2-2

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

用前缀和来计算。

提示 2-3

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

不妨将子数组的右端点固定，子数组左端点的范围是多少？

对于多个不同的右端点，其对应的左端点的范围是否均相同？

提示 2-4

设子数组左端点为 l，右端点为 r，当前枚举的元素下标为 i，那么有 L\le l\le i \le r\le R。

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

s[r+1]-s[l]

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

\begin{aligned}
&\sum_{r=i+1}^{R+1}\sum_{l=L}^{i} (s[r]-s[l]) \
=&\left(\sum_{r=i+1}^{R+1}\sum_{l=L}^{i} s[r]\right)-\left(\sum_{r=i+1}^{R+1}\sum_{l=L}^{i} s[l]\right) \
=&(i-L+1)\cdot \sum_{r=i+1}^{R+1}s[r] -(R-i+1)\cdot \sum_{l=L}^{i} s[l]
\end{aligned}

因此我们还需要计算出前缀和 s 的前缀和 ss，其中 ss}[i]=\sum\limits_{j=0}^{i-1}s[j]，上式即为

(i-L+1)\cdot (\textit{ss}[R+2]-\textit{ss}[i+1]) - (R-i+1)\cdot (\textit{ss}[i+1]-\textit{ss}[L])

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

[sol1-Python3]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def totalStrength(self, strength: List[int]) -> int:
        n = len(strength)
        # left[i] 为左侧严格小于 strength[i] 的最近元素位置（不存在时为 -1）
        # right[i] 为右侧小于等于 strength[i] 的最近元素位置（不存在时为 n）
        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  # [l, r]  左闭右闭
            tot = (i - l + 1) * (ss[r + 2] - ss[i + 1]) - (r - i + 1) * (ss[i + 1] - ss[l])
            ans += v * tot  # 累加贡献
        return ans % (10 ** 9 + 7)
[sol1-Java]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class Solution {
    public int totalStrength(int[] strength) {
        final var mod = (int) 1e9 + 7;

        var n = strength.length;
        var left = new int[n];  // left[i] 为左侧严格小于 strength[i] 的最近元素位置（不存在时为 -1）
        var right = new int[n]; // right[i] 为右侧小于等于 strength[i] 的最近元素位置（不存在时为 n）
        Arrays.fill(right, n);
        var st = new ArrayDeque<Integer>();
        st.push(-1); // 哨兵
        for (var i = 0; i < n; i++) {
            while (st.size() > 1 && strength[st.peek()] >= strength[i])
                right[st.pop()] = i;
            left[i] = st.peek();
            st.push(i);
        }

        var s = 0L; // 前缀和
        var ss = new int[n + 2]; // 前缀和的前缀和
        for (var i = 1; i <= n; ++i) {
            s += strength[i - 1];
            ss[i + 1] = (int) ((ss[i] + s) % mod); // 注意取模后，下面计算两个 ss 相减，结果可能为负
        }

        var ans = 0L;
        for (var i = 0; i < n; ++i) {
            int l = left[i] + 1, r = right[i] - 1; // [l,r] 左闭右闭
            var tot = ((long) (i - l + 1) * (ss[r + 2] - ss[i + 1]) - (long) (r - i + 1) * (ss[i + 1] - ss[l])) % mod;
            ans = (ans + strength[i] * tot) % mod; // 累加贡献
        }
        return (int) (ans + mod) % mod; // 防止算出负数
    }
}
[sol1-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
public:
    int totalStrength(vector<int> &strength) {
        const int mod = 1e9 + 7;

        int n = strength.size();
        vector<int> left(n, -1); // left[i] 为左侧严格小于 strength[i] 的最近元素位置（不存在时为 -1）
        vector<int> right(n, n); // right[i] 为右侧小于等于 strength[i] 的最近元素位置（不存在时为 n）
        stack<int> st;
        for (int i = 0; i < n; ++i) {
            while (!st.empty() && strength[st.top()] >= strength[i]) {
                right[st.top()] = i;
                st.pop();
            }
            if (!st.empty()) left[i] = st.top();
            st.push(i);
        }

        long s = 0L; // 前缀和
        vector<int> ss(n + 2); // 前缀和的前缀和
        for (int i = 1; i <= n; ++i) {
            s += strength[i - 1];
            ss[i + 1] = (ss[i] + s) % mod; // 注意取模后，下面计算两个 ss 相减，结果可能为负
        }

        int ans = 0;
        for (int i = 0; i < n; ++i) {
            long l = left[i] + 1, r = right[i] - 1; // [l,r] 左闭右闭
            long tot = ((i - l + 1) * (ss[r + 2] - ss[i + 1]) - (r - i + 1) * (ss[i + 1] - ss[l])) % mod;
            ans = (ans + strength[i] * tot) % mod; // 累加贡献
        }
        return (ans + mod) % mod; // 防止算出负数
    }
};
[sol1-Go]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
func totalStrength(strength []int) (ans int) {
	const mod int = 1e9 + 7

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

	s := 0 // 前缀和
	ss := make([]int, n+2) // 前缀和的前缀和
	for i, v := range strength {
		s += v
		ss[i+2] = (ss[i+1] + s) % mod // 注意取模后，下面计算两个 ss 相减，结果可能为负
	}
	for i, v := range strength {
		l, r := left[i]+1, right[i]-1 // [l,r] 左闭右闭
		tot := ((i-l+1)*(ss[r+2]-ss[i+1]) - (r-i+1)*(ss[i+1]-ss[l])) % mod
		ans = (ans + v*tot) % mod // 累加贡献
	}
	return (ans + mod) % mod // 防止算出负数
}

复杂度分析

  • 时间复杂度：O(n)。
  • 空间复杂度：O(n)。

相似题目

 Comments
On this page
2281-巫师的总力量和