class Solution {
public:
int totalStrength(vector<int>& strength) {
}
};
2281. 巫师的总力量和
作为国王的统治者，你有一支巫师军队听你指挥。
给你一个下标从 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 <= 1051 <= strength[i] <= 109原站题解
cpp 解法, 执行用时: 196 ms, 内存消耗: 89.8 MB, 提交时间: 2023-06-29 10:40:44
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; // 防止算出负数
}
};
golang 解法, 执行用时: 172 ms, 内存消耗: 9.9 MB, 提交时间: 2023-06-29 10:40:20
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 // 防止算出负数
}
java 解法, 执行用时: 35 ms, 内存消耗: 57.3 MB, 提交时间: 2023-06-29 10:40:01
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; // 防止算出负数
}
}
python3 解法, 执行用时: 520 ms, 内存消耗: 35.6 MB, 提交时间: 2023-06-29 10:39:34
# 单调栈 + 前缀和的前缀和
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)