Sum of Total Strength of Wizards
Problem
As the ruler of a kingdom, you have an army of wizards at your command.
You are given a 0-indexed integer array strength, where strength[i]
denotes the strength of the ith wizard. For a contiguous group of wizards (i.e. the wizards' strengths form a subarray of strength), the
total strength is defined as the product of the following two values:
- The strength of the weakest wizard in the group.
- The total of all the individual strengths of the wizards in the group.
Return thesum of the total strengths of all contiguous groups of wizards. Since the answer may be very large, return it modulo 10^9 + 7.
A subarray is a contiguous non-empty sequence of elements within an array.
Examples
Example 1
Input: strength = [1,3,1,2]
Output: 44
Explanation: The following are all the contiguous groups of wizards:
- [1] from [_**1**_ ,3,1,2] has a total strength of min([1]) * sum([1]) = 1 * 1 = 1
- [3] from [1,_**3**_ ,1,2] has a total strength of min([3]) * sum([3]) = 3 * 3 = 9
- [1] from [1,3,_**1**_ ,2] has a total strength of min([1]) * sum([1]) = 1 * 1 = 1
- [2] from [1,3,1,_**2**_] has a total strength of min([2]) * sum([2]) = 2 * 2 = 4
- [1,3] from [_**1,3**_ ,1,2] has a total strength of min([1,3]) * sum([1,3]) = 1 * 4 = 4
- [3,1] from [1,_**3,1**_ ,2] has a total strength of min([3,1]) * sum([3,1]) = 1 * 4 = 4
- [1,2] from [1,3,_**1,2**_] has a total strength of min([1,2]) * sum([1,2]) = 1 * 3 = 3
- [1,3,1] from [_**1,3,1**_ ,2] has a total strength of min([1,3,1]) * sum([1,3,1]) = 1 * 5 = 5
- [3,1,2] from [1,_**3,1,2**_] has a total strength of min([3,1,2]) * sum([3,1,2]) = 1 * 6 = 6
- [1,3,1,2] from [_**1,3,1,2**_] has a total strength of min([1,3,1,2]) * sum([1,3,1,2]) = 1 * 7 = 7
The sum of all the total strengths is 1 + 9 + 1 + 4 + 4 + 4 + 3 + 5 + 6 + 7 = 44.
Example 2
Input: strength = [5,4,6]
Output: 213
Explanation: The following are all the contiguous groups of wizards:
- [5] from [_**5**_ ,4,6] has a total strength of min([5]) * sum([5]) = 5 * 5 = 25
- [4] from [5,_**4**_ ,6] has a total strength of min([4]) * sum([4]) = 4 * 4 = 16
- [6] from [5,4,_**6**_] has a total strength of min([6]) * sum([6]) = 6 * 6 = 36
- [5,4] from [_**5,4**_ ,6] has a total strength of min([5,4]) * sum([5,4]) = 4 * 9 = 36
- [4,6] from [5,_**4,6**_] has a total strength of min([4,6]) * sum([4,6]) = 4 * 10 = 40
- [5,4,6] from [_**5,4,6**_] has a total strength of min([5,4,6]) * sum([5,4,6]) = 4 * 15 = 60
The sum of all the total strengths is 25 + 16 + 36 + 36 + 40 + 60 = 213.
Constraints
1 <= strength.length <= 10^51 <= strength[i] <= 10^9
Solution
Method 1 - Monotonic Stack + Prefix Sums (Contribution by Minimum)
Intuition
For each element in strength we compute its contribution as the minimum across all subarrays where it is the minimum; summing these contributions gives the final answer. We use a monotonic stack to find the nearest smaller element to the left (left[i]) and the nearest smaller-or-equal to the right (right[i]) for every index i. To get the sum of subarray sums efficiently we maintain prefix sums pre and prefix-of-prefix sums ppre, then compute each element's total contribution using those arrays and multiply by the element value. Use MOD for modulo arithmetic and accumulate into ans.
Approach
- Precompute
prewherepre[k]= sum ofstrength[0..k-1]moduloMOD. - Precompute
pprewhereppre[k]= sum ofpre[0..k-1]moduloMOD(prefix of prefix sums). - Use a monotonic increasing stack to compute
left[i]= index of previous smaller element for eachi(or-1if none). - Use a monotonic increasing stack (with >= for tie-breaking) from right to left to compute
right[i]= index of next smaller-or-equal element for eachi(ornif none). - For each index
i, letl = left[i]andr = right[i]. Usingppreandprewe compute the sum of all subarray sums whereiis the minimum astotal = ((ppre[r+1] - ppre[i+1]) * (i-l) - (ppre[i+1] - ppre[l+1]) * (r-i)) % MOD(adjusted to be positive). Multiplytotalbystrength[i]and add toansmoduloMOD. - Return
ans.
Code
C++
class Solution {
public:
int totalStrength(vector<int>& strength) {
const int MOD = 1e9+7;
int n = strength.size();
vector<long> pre(n+2), ppre(n+3);
for (int i = 0; i < n; ++i) pre[i+1] = (pre[i] + strength[i]) % MOD;
for (int i = 0; i <= n+1; ++i) ppre[i+1] = (ppre[i] + pre[i]) % MOD;
vector<int> left(n), right(n);
stack<int> st;
for (int i = 0; i < n; ++i) {
while (!st.empty() && strength[st.top()] > strength[i]) st.pop();
left[i] = st.empty() ? -1 : st.top();
st.push(i);
}
while (!st.empty()) st.pop();
for (int i = n-1; i >= 0; --i) {
while (!st.empty() && strength[st.top()] >= strength[i]) st.pop();
right[i] = st.empty() ? n : st.top();
st.push(i);
}
long ans = 0;
for (int i = 0; i < n; ++i) {
int l = left[i], r = right[i];
long total = ((ppre[r+1] - ppre[i+1] + MOD) % MOD * (i-l) % MOD
- (ppre[i+1] - ppre[l+1] + MOD) % MOD * (r-i) % MOD + MOD) % MOD;
ans = (ans + strength[i] * total % MOD) % MOD;
}
return (int)ans;
}
};
Java
import java.util.*;
class Solution {
public int totalStrength(int[] strength) {
int MOD = 1_000_000_007, n = strength.length;
long[] pre = new long[n+2], ppre = new long[n+3];
for (int i = 0; i < n; ++i) pre[i+1] = (pre[i] + strength[i]) % MOD;
for (int i = 0; i <= n+1; ++i) ppre[i+1] = (ppre[i] + pre[i]) % MOD;
int[] left = new int[n], right = new int[n];
Deque<Integer> st = new ArrayDeque<>();
for (int i = 0; i < n; ++i) {
while (!st.isEmpty() && strength[st.peek()] > strength[i]) st.pop();
left[i] = st.isEmpty() ? -1 : st.peek();
st.push(i);
}
st.clear();
for (int i = n-1; i >= 0; --i) {
while (!st.isEmpty() && strength[st.peek()] >= strength[i]) st.pop();
right[i] = st.isEmpty() ? n : st.peek();
st.push(i);
}
long ans = 0;
for (int i = 0; i < n; ++i) {
int l = left[i], r = right[i];
long total = ((ppre[r+1] - ppre[i+1] + MOD) % MOD * (i-l) % MOD
- (ppre[i+1] - ppre[l+1] + MOD) % MOD * (r-i) % MOD + MOD) % MOD;
ans = (ans + strength[i] * total % MOD) % MOD;
}
return (int)ans;
}
}
Python
class Solution:
def totalStrength(self, strength: list[int]) -> int:
MOD = 10**9 + 7
n = len(strength)
pre = [0] * (n+2)
ppre = [0] * (n+3)
for i in range(n):
pre[i+1] = (pre[i] + strength[i]) % MOD
for i in range(n+2):
ppre[i+1] = (ppre[i] + pre[i]) % MOD
left = [-1] * n
right = [n] * n
st = []
for i in range(n):
while st and strength[st[-1]] > strength[i]:
st.pop()
left[i] = st[-1] if st else -1
st.append(i)
st.clear()
for i in range(n-1, -1, -1):
while st and strength[st[-1]] >= strength[i]:
st.pop()
right[i] = st[-1] if st else n
st.append(i)
ans = 0
for i in range(n):
l, r = left[i], right[i]
total = ((ppre[r+1] - ppre[i+1]) * (i-l) - (ppre[i+1] - ppre[l+1]) * (r-i)) % MOD
ans = (ans + strength[i] * total) % MOD
return ans % MOD
Rust
impl Solution {
pub fn total_strength(strength: Vec<i32>) -> i32 {
const MOD: i64 = 1_000_000_007;
let n = strength.len();
let mut pre = vec![0i64; n+2];
let mut ppre = vec![0i64; n+3];
for i in 0..n { pre[i+1] = (pre[i] + strength[i] as i64) % MOD; }
for i in 0..=n+1 { ppre[i+1] = (ppre[i] + pre[i]) % MOD; }
let mut left = vec![-1; n];
let mut right = vec![n as i32; n];
let mut st = vec![];
for i in 0..n {
while let Some(&j) = st.last() {
if strength[j] > strength[i] { st.pop(); } else { break; }
}
left[i] = *st.last().unwrap_or(&-1) as i32;
st.push(i);
}
st.clear();
for i in (0..n).rev() {
while let Some(&j) = st.last() {
if strength[j] >= strength[i] { st.pop(); } else { break; }
}
right[i] = *st.last().unwrap_or(&(n as i32)) as i32;
st.push(i);
}
let mut ans = 0i64;
for i in 0..n {
let l = left[i] as usize;
let r = right[i] as usize;
let total = ((ppre[r+1] - ppre[i+1] + MOD) % MOD * (i as i64 - l as i64) % MOD
- (ppre[i+1] - ppre[l+1] + MOD) % MOD * (r as i64 - i as i64) % MOD + MOD) % MOD;
ans = (ans + strength[i] as i64 * total % MOD) % MOD;
}
ans as i32
}
}
TypeScript
function totalStrength(strength: number[]): number {
const MOD = 1e9 + 7;
const n = strength.length;
const pre = new Array(n+2).fill(0);
const ppre = new Array(n+3).fill(0);
for (let i = 0; i < n; ++i) pre[i+1] = (pre[i] + strength[i]) % MOD;
for (let i = 0; i <= n+1; ++i) ppre[i+1] = (ppre[i] + pre[i]) % MOD;
const left = new Array(n).fill(-1);
const right = new Array(n).fill(n);
const st: number[] = [];
for (let i = 0; i < n; ++i) {
while (st.length && strength[st[st.length-1]] > strength[i]) st.pop();
left[i] = st.length ? st[st.length-1] : -1;
st.push(i);
}
st.length = 0;
for (let i = n-1; i >= 0; --i) {
while (st.length && strength[st[st.length-1]] >= strength[i]) st.pop();
right[i] = st.length ? st[st.length-1] : n;
st.push(i);
}
let ans = 0;
for (let i = 0; i < n; ++i) {
const l = left[i], r = right[i];
const total = ((ppre[r+1] - ppre[i+1] + MOD) % MOD * (i-l) % MOD
- (ppre[i+1] - ppre[l+1] + MOD) % MOD * (r-i) % MOD + MOD) % MOD;
ans = (ans + strength[i] * total % MOD) % MOD;
}
return ans;
}
Complexity
- ⏰ Time complexity:
O(n) - 🧺 Space complexity:
O(n)