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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


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^5
1 <= strength[i] <= 10^9

Solution

Approach

We use a monotonic stack to find, for each element, the range where it is the minimum. We use prefix sums and prefix of prefix sums to efficiently compute the sum of all subarrays where each element is the minimum. The answer is the sum over all elements of their contribution as the minimum in their range.

Code

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


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

 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


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

 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


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

 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
35
36
37


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

 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


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;
}

Explanation

We use a monotonic stack to find the range where each element is the minimum, and prefix sums to compute the sum of all subarrays efficiently. Each element’s contribution is calculated and summed.

Complexity

⏰ Time complexity: O(n)
🧺 Space complexity: O(n)

Problem#

Examples#

Example 1#

Example 2#

Constraints#

Solution#

Approach#

Code#

C++#

Java#

Python#

Rust#

TypeScript#

Explanation#

Complexity#