Sum of Total Strength of Wizards

Hard

#2079Time: O(n^2), where n is the number of wizards. We have two nested loops to iterate through all `n * (n + 1) / 2` contiguous subarrays.Space: O(1), as we only use a constant amount of extra space for variables to store the current sum, minimum, and total strength.

Patterns

Prefix Sum

Data structures

Array Stack Monotonic Stack

Prompt

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 i^th 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 10⁹ + 7.

A subarray is a contiguous non-empty sequence of elements within an array.

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⁵
1 <= strength[i] <= 10⁹

Approaches

2 approaches with complexity analysis and trade-offs.

This approach involves iterating through every possible contiguous subarray. For each subarray, we calculate its sum and find its minimum element. The product of these two values gives the total strength of that subarray. We sum these strengths up for all subarrays to get the final answer. While simple, this method is computationally expensive.

Algorithm

Initialize a variable totalStrengthSum to 0 and MOD = 1_000_000_007.
Use a nested loop structure. The outer loop with index i from 0 to n-1 will define the starting point of the subarrays.
The inner loop with index j from i to n-1 will define the ending point of the subarrays.
Inside the outer loop, initialize currentSum = 0 and minVal = infinity.
Inside the inner loop, for each element strength[j], update the currentSum by adding strength[j] and update minVal by taking the minimum of the current minVal and strength[j].
Calculate the total strength of the current subarray strength[i..j] as (minVal * currentSum) % MOD.
Add this result to the totalStrengthSum, ensuring the sum also stays within the modulo.
After both loops complete, totalStrengthSum will hold the final result.

Walkthrough

The brute-force method systematically generates all contiguous subarrays. We can optimize the naive O(n^3) approach to O(n^2) by observing that when we extend a subarray strength[i..j] to strength[i..j+1], the sum and minimum can be updated in constant time. We use two nested loops: the outer loop fixes the start of the subarray, and the inner loop extends the end. For each subarray formed, we calculate its strength and add it to a running total. All calculations involving large numbers are performed under modulo 10^9 + 7 to prevent overflow.

1class Solution {2    public int totalStrength(int[] strength) {3        long totalStrengthSum = 0;4        int n = strength.length;5        long MOD = 1_000_000_007;6 7        for (int i = 0; i < n; i++) {8            long currentSum = 0;9            long minVal = Long.MAX_VALUE;10            for (int j = i; j < n; j++) {11                currentSum += strength[j];12                minVal = Math.min(minVal, strength[j]);13                14                long term1 = minVal % MOD;15                long term2 = currentSum % MOD;16                long currentTotalStrength = (term1 * term2) % MOD;17                18                totalStrengthSum = (totalStrengthSum + currentTotalStrength) % MOD;19            }20        }21        return (int) totalStrengthSum;22    }23}

Complexity

Time

O(n^2), where n is the number of wizards. We have two nested loops to iterate through all `n * (n + 1) / 2` contiguous subarrays.

Space

O(1), as we only use a constant amount of extra space for variables to store the current sum, minimum, and total strength.

Trade-offs

Pros

Simple to understand and straightforward to implement.

Cons

Highly inefficient for the given constraints (n <= 10^5), resulting in a Time Limit Exceeded (TLE) error on most platforms.

This optimal approach reframes the problem from iterating over subarrays to calculating the contribution of each element strength[i] to the final answer. For each element, we identify all subarrays where it is the minimum. The sum of these subarrays' sums is then calculated efficiently using a clever combination of a monotonic stack and prefix sums of prefix sums. This reduces the overall time complexity to linear.

Algorithm

Define MOD = 1_000_000_007.
For each index i, find left[i] (the index of the first element to the left of i that is strictly smaller than strength[i]) and right[i] (the index of the first element to the right of i that is smaller than or equal to strength[i]). This is done efficiently in O(n) using a monotonic stack. Using strict inequality for one side and non-strict for the other handles duplicate values correctly.
Compute the prefix sum array P of strength, where P[k] = sum(strength[0...k-1]). This takes O(n).
Compute the prefix sum of the prefix sum array, PP, where PP[k] = sum(P[0...k-1]). This also takes O(n).
Initialize a variable totalStrength = 0.
Iterate from i = 0 to n-1: a. Get the precomputed boundaries l = left[i] and r = right[i]. b. The total sum of sums for all subarrays where strength[i] is the minimum can be calculated using a formula derived from P and PP: sum_of_sums = (i - l) * (PP[r+1] - PP[i+1]) - (r - i) * (PP[i+1] - PP[l+1]). c. Calculate this sum_of_sums using modulo arithmetic, being careful with subtractions which might yield negative results. d. The contribution of strength[i] is (strength[i] * sum_of_sums) % MOD. e. Add this contribution to totalStrength.
Return the final totalStrength.

Walkthrough

Instead of summing min * sum for each subarray, we can change our perspective: for each element strength[i], what is its total contribution to the final sum? The contribution is strength[i] multiplied by the sum of sums of all subarrays for which strength[i] is the minimum element.

1. Finding Ranges with Monotonic Stack: We first need to find the boundaries for each strength[i] where it acts as the minimum. We can find left[i], the index of the previous smaller element, and right[i], the index of the next smaller or equal element. Any subarray strength[j..k] with left[i] < j <= i <= k < right[i] will have strength[i] as its minimum. This can be done in O(n) time using a monotonic stack.

2. Efficiently Calculating Sum of Subarray Sums: Calculating the sum of sum(strength[j..k]) for all valid j and k for a given i naively would be too slow. We can precompute prefix sums of the strength array (let's call it P) and then prefix sums of that prefix sum array (let's call it PP). These structures allow us to calculate the sum of sums over a range in O(1) time. The final formula for the sum of sums for element i is (i - left[i]) * (PP[right[i]+1] - PP[i+1]) - (right[i] - i) * (PP[i+1] - PP[left[i]+1]).

By combining these techniques, we can calculate the contribution of each element in O(1) time (after O(n) precomputation), leading to an overall O(n) solution.

1import java.util.Stack;2 3class Solution {4    public int totalStrength(int[] strength) {5        int n = strength.length;6        long MOD = 1_000_000_007;7 8        // left[i]: index of the first element to the left of i that is strictly smaller than strength[i]9        int[] left = new int[n];10        Stack<Integer> stack = new Stack<>();11        for (int i = 0; i < n; i++) {12            while (!stack.isEmpty() && strength[stack.peek()] >= strength[i]) {13                stack.pop();14            }15            left[i] = stack.isEmpty() ? -1 : stack.peek();16            stack.push(i);17        }18 19        // right[i]: index of the first element to the right of i that is smaller than or equal to strength[i]20        int[] right = new int[n];21        stack.clear();22        for (int i = n - 1; i >= 0; i--) {23            while (!stack.isEmpty() && strength[stack.peek()] > strength[i]) {24                stack.pop();25            }26            right[i] = stack.isEmpty() ? n : stack.peek();27            stack.push(i);28        }29 30        // prefixSum[i] = sum(strength[0]...strength[i-1])31        long[] prefixSum = new long[n + 1];32        for (int i = 0; i < n; i++) {33            prefixSum[i + 1] = (prefixSum[i] + strength[i]) % MOD;34        }35 36        // prefixSumOfPrefixSum[i] = sum(prefixSum[0]...prefixSum[i-1])37        long[] prefixSumOfPrefixSum = new long[n + 2];38        for (int i = 0; i < n + 1; i++) {39            prefixSumOfPrefixSum[i + 1] = (prefixSumOfPrefixSum[i] + prefixSum[i]) % MOD;40        }41 42        long totalStrength = 0;43        for (int i = 0; i < n; i++) {44            int l = left[i];45            int r = right[i];46 47            long numLeft = i - l;48            long numRight = r - i;49 50            long sumLeftPart = (prefixSumOfPrefixSum[i + 1] - prefixSumOfPrefixSum[l + 1] + MOD) % MOD;51            long sumRightPart = (prefixSumOfPrefixSum[r + 1] - prefixSumOfPrefixSum[i + 1] + MOD) % MOD;52            53            long term1 = (numLeft * sumRightPart) % MOD;54            long term2 = (numRight * sumLeftPart) % MOD;55 56            long sumOfSums = (term1 - term2 + MOD) % MOD;57            58            totalStrength = (totalStrength + (strength[i] * sumOfSums) % MOD) % MOD;59        }60 61        return (int) totalStrength;62    }63}

Complexity

Time

O(n), where n is the number of wizards. Each major step—computing `left`/`right` arrays with a monotonic stack, computing prefix sums, and the final loop to sum contributions—takes linear time.

Space

O(n), where n is the number of wizards. We use several arrays (`left`, `right`, `prefixSum`, `prefixSumOfPrefixSum`) and a stack, all of which require space proportional to the input size.

Trade-offs

Pros

Optimal time complexity, making it very fast and capable of handling large inputs.
Demonstrates a powerful problem-solving pattern by calculating the contribution of each element, which is applicable to a class of similar problems.

Cons

The logic, especially the derivation of the formula for the sum of sums, is complex and non-intuitive.
Implementation is prone to off-by-one errors due to complex indexing and requires careful handling of modulo arithmetic.

Solutions

Solution

1class Solution {2public3  int totalStrength(int[] strength) {4    int n = strength.length;5    int[] left = new int[n];6    int[] right = new int[n];7    Arrays.fill(left, -1);8    Arrays.fill(right, n);9    Deque<Integer> stk = new ArrayDeque<>();10    for (int i = 0; i < n; ++i) {11      while (!stk.isEmpty() && strength[stk.peek()] >= strength[i]) {12        stk.pop();13      }14      if (!stk.isEmpty()) {15        left[i] = stk.peek();16      }17      stk.push(i);18    }19    stk.clear();20    for (int i = n - 1; i >= 0; --i) {21      while (!stk.isEmpty() && strength[stk.peek()] > strength[i]) {22        stk.pop();23      }24      if (!stk.isEmpty()) {25        right[i] = stk.peek();26      }27      stk.push(i);28    }29    int mod = (int)1 e9 + 7;30    int[] s = new int[n + 1];31    for (int i = 0; i < n; ++i) {32      s[i + 1] = (s[i] + strength[i]) % mod;33    }34    int[] ss = new int[n + 2];35    for (int i = 0; i < n + 1; ++i) {36      ss[i + 1] = (ss[i] + s[i]) % mod;37    }38    long ans = 0;39    for (int i = 0; i < n; ++i) {40      int v = strength[i];41      int l = left[i] + 1, r = right[i] - 1;42      long a = (long)(i - l + 1) * (ss[r + 2] - ss[i + 1]);43      long b = (long)(r - i + 1) * (ss[i + 1] - ss[l]);44      ans = (ans + v * ((a - b) % mod)) % mod;45    }46    return (int)(ans + mod) % mod;47  }48}

1class Solution {2public:3  int totalStrength(vector<int> &strength) {4    int n = strength.size();5    vector<int> left(n, -1);6    vector<int> right(n, n);7    stack<int> stk;8    for (int i = 0; i < n; ++i) {9      while (!stk.empty() && strength[stk.top()] >= strength[i])10        stk.pop();11      if (!stk.empty())12        left[i] = stk.top();13      stk.push(i);14    }15    stk = stack<int>();16    for (int i = n - 1; i >= 0; --i) {17      while (!stk.empty() && strength[stk.top()] > strength[i])18        stk.pop();19      if (!stk.empty())20        right[i] = stk.top();21      stk.push(i);22    }23    int mod = 1e9 + 7;24    vector<int> s(n + 1);25    for (int i = 0; i < n; ++i)26      s[i + 1] = (s[i] + strength[i]) % mod;27    vector<int> ss(n + 2);28    for (int i = 0; i < n + 1; ++i)29      ss[i + 1] = (ss[i] + s[i]) % mod;30    int ans = 0;31    for (int i = 0; i < n; ++i) {32      int v = strength[i];33      int l = left[i] + 1, r = right[i] - 1;34      long a = (long)(i - l + 1) * (ss[r + 2] - ss[i + 1]);35      long b = (long)(r - i + 1) * (ss[i + 1] - ss[l]);36      ans = (ans + v * ((a - b) % mod)) % mod;37    }38    return (int)(ans + mod) % mod;39  }40};

1class Solution:2    def totalStrength(self, strength: List[int]) -> int: n = len(strength) left = [- 1] * n right = [n] * n stk = [] for i, v in enumerate(strength): while stk and strength[stk[- 1]] >= v: stk . pop() if stk: left[i] = stk[- 1] stk . append(i) stk = [] for i in range(n - 1, - 1, - 1): while stk and strength[stk[- 1]] > strength[i]: stk . pop() if stk: right[i] = stk[- 1] stk . append(i) ss = list(accumulate(list(accumulate(strength, initial=0)), initial=0)) mod = int(1e9) + 7 ans = 0 for i, v in enumerate(strength): l, r = left[i] + 1, right[i] - 1 a = (ss[r + 2] - ss[i + 1]) * (i - l + 1) b = (ss[i + 1] - ss[l]) * (r - i + 1) ans = (ans + (a - b) * v) % mod return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem