Sum of total strength of wizards

Array Stack Monotonic Stack Prefix Sum

29.2% acceptance

Feb 25, 2026

1267

111

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 the sum of the total strengths of all contiguous groups of wizards. Since the answer may be very large, return it modulo 109 + 7. A subarray is a contiguous non-empty sequence of elements within an array.

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn total_strength(strength: Vec<i32>) -> i32 {
    const MOD: i64 = 1_000_000_007;
    let n = strength.len();

    // prefix sums P[i] = sum of strength[0..i]
    let mut p = vec![0i64; n + 1];
    for i in 0..n { p[i+1] = (p[i] + strength[i] as i64) % MOD; }

    // prefix sums of P: Q[i] = sum of P[0..i]
    let mut q = vec![0i64; n + 2];
    for i in 0..=n { q[i+1] = (q[i] + p[i]) % MOD; }

    // Monotonic stack to find left[i] (largest j < i with strength[j] < strength[i])
    // and right[i] (smallest j > i with strength[j] <= strength[i])
    let mut left = vec![-1i64; n];
    let mut right = vec![n as i64; n];
    let mut stack: Vec<usize> = Vec::new();

    // Left boundary (strict less)
    for i in 0..n {
      while !stack.is_empty() && strength[*stack.last().unwrap()] >= strength[i] {
        stack.pop();
      }
      if let Some(&top) = stack.last() { left[i] = top as i64; }
      stack.push(i);
    }

    // Right boundary (less-or-equal)
    stack.clear();
    for i in (0..n).rev() {
      while !stack.is_empty() && strength[*stack.last().unwrap()] > strength[i] {
        stack.pop();
      }
      if let Some(&top) = stack.last() { right[i] = top as i64; }
      stack.push(i);
    }

    let mut ans = 0i64;
    for i in 0..n {
      let l = left[i];   // left[i] is the last pos with strictly smaller value, or -1
      let r = right[i];  // right[i] is the first pos with smaller-or-equal value, or n
      let left_cnt = (i as i64 - l) % MOD;
      let right_cnt = (r - i as i64) % MOD;
      // sum_right_P = Q[r+1] - Q[i+1]
      let sum_rp = (q[r as usize + 1] - q[i + 1] + MOD) % MOD;
      // sum_left_P = Q[i+1] - Q[l+1]
      let sum_lp = (q[i + 1] - q[(l + 1) as usize] + MOD) % MOD;
      let contrib = strength[i] as i64 % MOD
        * ((left_cnt * sum_rp % MOD - right_cnt * sum_lp % MOD + MOD) % MOD) % MOD;
      ans = (ans + contrib) % MOD;
    }
    ans as i32
  }
}