Sum of Total Strength of Wizards

hard array monotonic stack prefix sum

Problem

Given a 0-indexed array strength, the total strength of a contiguous group (subarray) is the product of its minimum element and the sum of all its elements. Return the sum of total strengths over every non-empty subarray, modulo 10⁹ + 7.

Inputstrength = [1, 3, 1, 2]

Output44

Across all 10 subarrays the totals min × sum are 1, 9, 1, 4, 4, 4, 3, 5, 6, 7, which add up to 44.

strength (comma-separated, positive)

def totalStrength(strength):
    MOD = 10**9 + 7
    n = len(strength)
    prefix = [0] * (n + 1)                 # prefix[i] = sum of strength[0..i-1]
    for i in range(n):
        prefix[i + 1] = prefix[i] + strength[i]
    pp = [0] * (n + 2)                      # prefix of prefix
    for i in range(n + 1):
        pp[i + 1] = pp[i] + prefix[i]
    left, right, st = [-1] * n, [n] * n, []
    for i in range(n):                     # previous strictly-smaller index
        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):         # next smaller-or-equal index
        while st and strength[st[-1]] > strength[i]:
            st.pop()
        right[i] = st[-1] if st else n
        st.append(i)
    ans = 0
    for j in range(n):                     # j is the minimum of its span
        L, R = left[j] + 1, right[j] - 1
        sumR = pp[R + 2] - pp[j + 1]       # sum of prefix[r+1], r in [j..R]
        sumL = pp[j + 1] - pp[L]           # sum of prefix[l],   l in [L..j]
        contrib = strength[j] * ((j - L + 1) * sumR - (R - j + 1) * sumL)
        ans = (ans + contrib) % MOD
    return ans % MOD

function totalStrength(strength) {
  const MOD = 1000000007n;
  const n = strength.length;
  const prefix = new Array(n + 1).fill(0n);        // prefix[i] = sum 0..i-1
  for (let i = 0; i < n; i++) prefix[i + 1] = prefix[i] + BigInt(strength[i]);
  const pp = new Array(n + 2).fill(0n);            // prefix of prefix
  for (let i = 0; i <= n; i++) pp[i + 1] = pp[i] + prefix[i];
  const left = new Array(n).fill(-1), right = new Array(n).fill(n);
  let st = [];
  for (let i = 0; i < n; i++) {                     // previous strictly-smaller
    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 = [];
  for (let i = n - 1; i >= 0; i--) {                // next smaller-or-equal
    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 = 0n;
  for (let j = 0; j < n; j++) {                     // j is the minimum
    const L = left[j] + 1, R = right[j] - 1;
    const sumR = pp[R + 2] - pp[j + 1];
    const sumL = pp[j + 1] - pp[L];
    const c = BigInt(strength[j]) * (BigInt(j - L + 1) * sumR - BigInt(R - j + 1) * sumL);
    ans = (ans + c) % MOD;
  }
  return Number(((ans % MOD) + MOD) % MOD);
}

class Solution {
    public int totalStrength(int[] strength) {
        final long MOD = 1000000007L;
        int n = strength.length;
        long[] prefix = new long[n + 1];             // prefix[i] = sum 0..i-1
        for (int i = 0; i < n; i++) prefix[i + 1] = prefix[i] + strength[i];
        long[] pp = new long[n + 2];                 // prefix of prefix
        for (int i = 0; i <= n; i++) pp[i + 1] = (pp[i] + prefix[i] % MOD) % MOD;
        int[] left = new int[n], right = new int[n];
        Deque<Integer> st = new ArrayDeque<>();
        for (int i = 0; i < n; i++) {                 // previous strictly-smaller
            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--) {            // next smaller-or-equal
            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 j = 0; j < n; j++) {                 // j is the minimum
            long L = left[j] + 1, R = right[j] - 1;
            long sumR = (pp[(int)R + 2] - pp[j + 1] + MOD) % MOD;
            long sumL = (pp[j + 1] - pp[(int)L] + MOD) % MOD;
            long c = strength[j] % MOD * (((j - L + 1) % MOD * sumR
                     - (R - j + 1) % MOD * sumL) % MOD + MOD) % MOD;
            ans = (ans + c) % MOD;
        }
        return (int) (ans % MOD);
    }
}

int totalStrength(vector<int>& strength) {
    const long long MOD = 1000000007LL;
    int n = (int) strength.size();
    vector<long long> prefix(n + 1, 0);               // prefix[i] = sum 0..i-1
    for (int i = 0; i < n; i++) prefix[i + 1] = prefix[i] + strength[i];
    vector<long long> pp(n + 2, 0);                   // prefix of prefix
    for (int i = 0; i <= n; i++) pp[i + 1] = (pp[i] + prefix[i] % MOD) % MOD;
    vector<int> left(n, -1), right(n, n);
    stack<int> st;
    for (int i = 0; i < n; i++) {                     // previous strictly-smaller
        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--) {                // next smaller-or-equal
        while (!st.empty() && strength[st.top()] > strength[i]) st.pop();
        right[i] = st.empty() ? n : st.top();
        st.push(i);
    }
    long long ans = 0;
    for (int j = 0; j < n; j++) {                     // j is the minimum
        long long L = left[j] + 1, R = right[j] - 1;
        long long sumR = (pp[R + 2] - pp[j + 1] + MOD) % MOD;
        long long sumL = (pp[j + 1] - pp[L] + MOD) % MOD;
        long long c = strength[j] % MOD * (((j - L + 1) % MOD * sumR
                       - (R - j + 1) % MOD * sumL) % MOD + MOD) % MOD;
        ans = (ans + c) % MOD;
    }
    return (int) (ans % MOD);
}

Explanation

There are O(n²) subarrays, so we cannot enumerate them. Instead we ask, for each element strength[j]: across all subarrays in which strength[j] is the minimum, what is the total of every subarray's min × sum? Adding those contributions over all j gives the answer, because every subarray has exactly one element acting as its minimum.

The span in which j is the minimum stretches left until an element strictly smaller appears and right until an element smaller appears. To avoid double-counting when equal values repeat, we make one side strict and the other inclusive: a monotonic stack gives left[j] (previous strictly-smaller index) and right[j] (next smaller-or-equal index). Then a valid subarray starts at any l in [L, j] and ends at any r in [j, R], where L = left[j]+1 and R = right[j]-1.

For a fixed j we must add strength[j] · Σ (prefix[r+1] − prefix[l]) over all those l and r. Splitting the double sum, each prefix[r+1] is counted once per choice of l (there are j − L + 1 of them) and each prefix[l] is subtracted once per choice of r (there are R − j + 1 of them).

The remaining sums of consecutive prefix values are themselves answered in O(1) by a second-level prefix array pp (a prefix of the prefix sums): Σ prefix[r+1] for r ∈ [j, R] is pp[R+2] − pp[j+1], and Σ prefix[l] for l ∈ [L, j] is pp[j+1] − pp[L]. So each element's whole contribution is a constant number of array lookups.

Example: strength = [1, 3, 1, 2]. The element 3 is the minimum only of [3] (contribution 9). The first 1 covers subarrays ending no later than the second 1; the second 1 covers the rest, including the full array. Element 2 is the minimum of [2] (contribution 4). The contributions sum to 44.

Building the prefix arrays and both boundary passes are each linear, and the final accumulation visits every index once, so the algorithm runs in O(n) time. All arithmetic is kept modulo 10⁹ + 7 (using a 64-bit or big-integer intermediate so the subtraction never underflows).

Time: O(n) Space: O(n)