Sum of Total Strength of Wizards

hardOA Simulation: 2 problems in 90 min each round96 stepsJava · Python · C++

For every subarray, strength = min × sum; return the total mod 1e9+7. Count by the minimum: for each wizard i, a monotonic stack finds the range where it's the minimum. The sum of all subarray sums inside that range comes from a prefix-of-prefix-sums array in O(1).

▶ Watch the animated trace

Traced input: strength = [1, 3, 1, 2] → 44. Original problem: leetcode.com/problems/sum-of-total-strength-of-wizards/

Complexity

BestO(n)

AverageO(n)

WorstO(n)

SpaceO(n)

Two monotonic-stack passes and two prefix arrays.

Key idea

Each subarray is assigned to one minimum: the previous smaller element (strict) and next smaller-or-equal bound it, so ties aren't double-counted.
Subarrays with min at i: start in (L, i], end in [i, R).
Sum over those of subarray sums = (pp[R+1] − pp[i+1])·(i − L) − (pp[i+1] − pp[L+1])·(R − i), where pp is the prefix sum of prefix sums.
Multiply by strength[i] and add, all mod 1e9+7.

Common mistakes

Using strict comparisons on both sides double-counts equal minimums.
Negative values after subtracting under a modulus; add MOD back.
O(n²) enumeration times out at n = 10⁵.

Solution

Java

int totalStrength(int[] strength) {
    final long MOD = 1_000_000_007L;
    int n = strength.length;
    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[] prefix = new long[n + 1], pp = new long[n + 2];
    for (int i = 0; i < n; i++) prefix[i + 1] = (prefix[i] + strength[i]) % MOD;
    for (int i = 0; i <= n; i++) pp[i + 1] = (pp[i] + prefix[i]) % MOD;
    long ans = 0;
    for (int i = 0; i < n; i++) {
        int l = left[i], r = right[i];
        long pos = (pp[r + 1] - pp[i + 1] + MOD) % MOD * (i - l) % MOD;
        long neg = (pp[i + 1] - pp[l + 1] + MOD) % MOD * (r - i) % MOD;
        long sums = (pos - neg + MOD) % MOD;
        ans = (ans + strength[i] * sums) % MOD;
    }
    return (int) ans;
}

Python

def totalStrength(strength: List[int]) -> int:
    MOD = 10**9 + 7
    n = len(strength)
    left, right, st = [-1] * n, [n] * n, []
    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 = []
    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)
    prefix = [0] * (n + 1)
    for i in range(n): prefix[i + 1] = (prefix[i] + strength[i]) % MOD
    pp = [0] * (n + 2)
    for i in range(n + 1): pp[i + 1] = (pp[i] + prefix[i]) % MOD
    ans = 0
    for i in range(n):
        l, r = left[i], right[i]
        pos = (pp[r + 1] - pp[i + 1]) * (i - l)
        neg = (pp[i + 1] - pp[l + 1]) * (r - i)
        ans = (ans + strength[i] * (pos - neg)) % MOD
    return ans

C++

int totalStrength(vector<int>& strength) {
    const long long MOD = 1'000'000'007LL;
    int n = strength.size();
    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);
    }
    st = stack<int>();
    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);
    }
    vector<long long> prefix(n + 1), pp(n + 2);
    for (int i = 0; i < n; i++) prefix[i + 1] = (prefix[i] + strength[i]) % MOD;
    for (int i = 0; i <= n; i++) pp[i + 1] = (pp[i] + prefix[i]) % MOD;
    long long ans = 0;
    for (int i = 0; i < n; i++) {
        int l = left[i], r = right[i];
        long long pos = (pp[r + 1] - pp[i + 1] + MOD) % MOD * (i - l) % MOD;
        long long neg = (pp[i + 1] - pp[l + 1] + MOD) % MOD * (r - i) % MOD;
        long long sums = (pos - neg + MOD) % MOD;
        ans = (ans + strength[i] * sums) % MOD;
    }
    return ans;
}

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
LeetCode example 2 Six subarrays: 25 + 16 + 36 + 36 + 40 + 60.	`strength = [5, 4, 6]`	`213`
Single wizard 7 × 7.	`strength = [7]`	`49`
Equal strengths 4 + 4 + 2×4 = 16; ties are counted once each.	`strength = [2, 2]`	`16`