LeetCode 2281 - Sum of Total Strength of Wizards

The problem asks us to compute the sum of total strengths across every possible contiguous subarray of the input array strength.

leetcodehardarraystackmonotonic-stackprefix-sum

LeetCode Problem 2281

Difficulty: 🔴 Hard
Topics: Array, Stack, Monotonic Stack, Prefix Sum

Solution

Problem Understanding

The problem asks us to compute the sum of total strengths across every possible contiguous subarray of the input array strength.

For any chosen subarray, its total strength is defined as:

$$\text{minimum value in subarray} \times \text{sum of values in subarray}$$

We must evaluate this quantity for all non-empty contiguous subarrays, then add all results together.

For example, if:

strength = [1,3,1,2]

then valid subarrays include:

[1]
[3]
[1,3]
[3,1]
[1,3,1]
[1,3,1,2]

and many others. For each one, we compute:

minimum(subarray) * sum(subarray)

Finally, we sum all these values.

The challenge comes from the constraints:

1 <= strength.length <= 100000
1 <= strength[i] <= 1000000000

A length of 10^5 immediately rules out any approach that explicitly enumerates all subarrays.

There are:

$$O(n^2)$$

subarrays in an array of size n.

Even computing the sum and minimum efficiently per subarray still becomes too slow because:

$$O(n^2)$$

for n = 100000 is far beyond practical limits.

This means we need a solution close to:

O(n log n)

or ideally:

O(n)

Another important complication is handling duplicate values correctly. When multiple equal elements exist, we must carefully define which index is considered the minimum for overlapping ranges, otherwise we risk double counting.

The problem guarantees:

  • The array is non-empty.
  • All values are positive integers.
  • The result can become extremely large, so we must return it modulo:

$$10^9 + 7$$

Important edge cases include:

  • A single element array.
  • All equal values.
  • Strictly increasing arrays.
  • Strictly decreasing arrays.
  • Large repeated values.
  • Cases where equal minima appear multiple times.

These situations often expose mistakes in monotonic stack boundary handling.

Approaches

Brute Force Approach

A straightforward solution is to enumerate every possible subarray.

For each starting index i, we iterate through every ending index j. While extending the subarray, we maintain:

  • the running sum
  • the current minimum

Then we compute:

minimum * subarray_sum

and add it to the final answer.

Although this gives the correct result, it is far too slow.

There are:

$$O(n^2)$$

subarrays.

Even if we update the minimum and sum incrementally in constant time, we still must examine every subarray.

With:

n = 100000

this becomes infeasible.

Key Insight for the Optimal Solution

Instead of iterating over subarrays, we reverse the thinking.

Rather than asking:

For every subarray, what is its minimum?

we ask:

For every element, in which subarrays is it the minimum?

This transformation is the key insight.

Suppose strength[i] is the minimum element.

We want to count the contribution of this wizard across all subarrays where it serves as the minimum.

To do this, we determine:

  • the nearest smaller element on the left
  • the nearest smaller or equal element on the right

These boundaries define exactly which subarrays choose strength[i] as their minimum.

This is a classic monotonic stack pattern.

However, we also need:

minimum × subarray sum

The minimum is fixed for an index i, but the subarray sum changes.

We therefore use prefix sums of prefix sums to efficiently calculate the sum of all subarray sums inside a range.

This reduces the total complexity to linear time.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(n²) O(1) Enumerates all subarrays directly
Optimal O(n) O(n) Monotonic stack + double prefix sums

Algorithm Walkthrough

Step 1: Compute Previous Smaller Elements

We first compute the nearest index to the left where the value is strictly smaller.

We use a monotonic increasing stack.

For each index i:

  • While stack top is greater than or equal to current value, pop.
  • The remaining top becomes the left boundary.
  • If no smaller element exists, boundary becomes -1.

We use strictly smaller on the left to avoid duplicate counting.

Step 2: Compute Next Smaller or Equal Elements

Next, we compute the nearest index to the right where the value is smaller or equal.

Again, we use a monotonic increasing stack.

For each index moving right-to-left:

  • While stack top is strictly greater than current value, pop.
  • Remaining top becomes right boundary.
  • If none exists, boundary becomes n.

This asymmetry:

left = strictly smaller
right = smaller or equal

ensures duplicate minimums are counted exactly once.

Step 3: Build Prefix Sum Array

We build:

prefix[i]

where:

prefix[i]

stores the sum of the first i elements.

This allows:

sum(l:r)

to be computed in constant time.

Step 4: Build Prefix of Prefix Sums

We also build:

prefix_of_prefix

This lets us efficiently compute the sum of many range sums.

Without this step, we would still spend too much time aggregating subarray sums.

Step 5: Compute Contribution of Each Element

Suppose index i has value:

x = strength[i]

Let:

l = left[i]
r = right[i]

Then:

  • left choices = i - l
  • right choices = r - i

We compute:

positive_part

for right expansions and

negative_part

for left expansions.

The formula becomes:

$$\text{totalSubarraySum} = (\text{rightContribution} \times \text{leftCount}) - (\text{leftContribution} \times \text{rightCount})$$

Then multiply by:

strength[i]

because this wizard is the minimum.

Add result modulo:

$$10^9+7$$

Step 6: Return Final Answer

After processing every index, return:

answer % MOD

Why it works

For every index i, the monotonic stack uniquely identifies all subarrays where strength[i] is the minimum. No subarray is missed, and none are counted twice because of the asymmetric comparison rules.

The double prefix sums guarantee that the total sum of all candidate subarrays can be computed in constant time for each index. Since every element is pushed and popped at most once from the stack, the entire algorithm runs in linear time.

Python Solution

from typing import List

class Solution:
    def totalStrength(self, strength: List[int]) -> int:
        MOD = 10**9 + 7
        n = len(strength)

        left = [-1] * n
        right = [n] * n

        stack = []

        # Previous strictly smaller
        for i in range(n):
            while stack and strength[stack[-1]] >= strength[i]:
                stack.pop()

            left[i] = stack[-1] if stack else -1
            stack.append(i)

        stack.clear()

        # Next smaller or equal
        for i in range(n - 1, -1, -1):
            while stack and strength[stack[-1]] > strength[i]:
                stack.pop()

            right[i] = stack[-1] if stack else n
            stack.append(i)

        # Prefix sums
        prefix = [0] * (n + 1)
        for i in range(n):
            prefix[i + 1] = (prefix[i] + strength[i]) % MOD

        # Prefix of prefix sums
        prefix_prefix = [0] * (n + 2)
        for i in range(n + 1):
            prefix_prefix[i + 1] = (
                prefix_prefix[i] + prefix[i]
            ) % MOD

        result = 0

        for i in range(n):
            left_idx = left[i]
            right_idx = right[i]

            left_count = i - left_idx
            right_count = right_idx - i

            positive_sum = (
                prefix_prefix[right_idx + 1]
                - prefix_prefix[i + 1]
            ) % MOD

            negative_sum = (
                prefix_prefix[i + 1]
                - prefix_prefix[left_idx + 1]
            ) % MOD

            total_sum = (
                positive_sum * left_count
                - negative_sum * right_count
            ) % MOD

            contribution = (
                strength[i] * total_sum
            ) % MOD

            result = (result + contribution) % MOD

        return result

The implementation begins by computing the valid span where each wizard can act as the minimum. The left array stores the nearest strictly smaller element, while right stores the nearest smaller-or-equal element.

After boundary discovery, we build a normal prefix sum array so range sums can be queried in constant time.

The important optimization is prefix_prefix, which stores cumulative prefix sums. This enables fast aggregation of many subarray sums at once.

For each index, we calculate all subarray sums where that wizard is the minimum, then multiply by the wizard’s strength and add to the answer.

The modulo operation is applied throughout to avoid overflow and maintain correctness.

Go Solution

func totalStrength(strength []int) int {
	const MOD int64 = 1_000_000_007

	n := len(strength)

	left := make([]int, n)
	right := make([]int, n)

	for i := range right {
		right[i] = n
		left[i] = -1
	}

	stack := []int{}

	// Previous strictly smaller
	for i := 0; i < n; i++ {
		for len(stack) > 0 &&
			strength[stack[len(stack)-1]] >= strength[i] {
			stack = stack[:len(stack)-1]
		}

		if len(stack) > 0 {
			left[i] = stack[len(stack)-1]
		}

		stack = append(stack, i)
	}

	stack = []int{}

	// Next smaller or equal
	for i := n - 1; i >= 0; i-- {
		for len(stack) > 0 &&
			strength[stack[len(stack)-1]] > strength[i] {
			stack = stack[:len(stack)-1]
		}

		if len(stack) > 0 {
			right[i] = stack[len(stack)-1]
		}

		stack = append(stack, i)
	}

	// Prefix sums
	prefix := make([]int64, n+1)

	for i := 0; i < n; i++ {
		prefix[i+1] = (
			prefix[i] + int64(strength[i])
		) % MOD
	}

	// Prefix of prefix sums
	prefixPrefix := make([]int64, n+2)

	for i := 0; i <= n; i++ {
		prefixPrefix[i+1] = (
			prefixPrefix[i] + prefix[i]
		) % MOD
	}

	var result int64 = 0

	for i := 0; i < n; i++ {
		leftIdx := left[i]
		rightIdx := right[i]

		leftCount := int64(i - leftIdx)
		rightCount := int64(rightIdx - i)

		positiveSum := (
			prefixPrefix[rightIdx+1] -
				prefixPrefix[i+1] +
				MOD,
		) % MOD

		negativeSum := (
			prefixPrefix[i+1] -
				prefixPrefix[leftIdx+1] +
				MOD,
		) % MOD

		totalSum := (
			positiveSum*leftCount -
				negativeSum*rightCount,
		) % MOD

		if totalSum < 0 {
			totalSum += MOD
		}

		contribution := (
			int64(strength[i]) * totalSum,
		) % MOD

		result = (result + contribution) % MOD
	}

	return int(result)
}

The Go implementation follows the same algorithmic structure as Python, but uses int64 extensively to avoid overflow. Since strengths can be as large as 10^9, intermediate multiplication can exceed the range of standard int.

Slices are used as stacks, and modulo correction is applied whenever subtraction might produce a negative result.

Worked Examples

Example 1

strength = [1,3,1,2]

Boundary Computation

Index Value Left Smaller Right Smaller/Equal
0 1 -1 2
1 3 0 2
2 1 -1 4
3 2 2 4

Prefix Arrays

Prefix:

i Value
0 0
1 1
2 4
3 5
4 7

Prefix of Prefix:

i Value
0 0
1 0
2 1
3 5
4 10
5 17

Contributions

Index Value Contribution
0 1 17
1 3 9
2 1 14
3 2 4

Final answer:

17 + 9 + 14 + 4 = 44

Example 2

strength = [5,4,6]

Boundary Computation

Index Value Left Smaller Right Smaller/Equal
0 5 -1 1
1 4 -1 3
2 6 1 3

Contributions

Index Value Contribution
0 5 25
1 4 136
2 6 36

Final answer:

25 + 136 + 36 = 213

Complexity Analysis

Measure Complexity Explanation
Time O(n) Each element enters and leaves the stack at most once
Space O(n) Stack, boundaries, and prefix arrays

The linear runtime comes from the monotonic stack property. Each index is pushed and popped once, and all prefix sum lookups are constant time. No nested iteration over subarrays occurs.

Test Cases

sol = Solution()

assert sol.totalStrength([1]) == 1  # single element
assert sol.totalStrength([1, 3, 1, 2]) == 44  # example 1
assert sol.totalStrength([5, 4, 6]) == 213  # example 2

assert sol.totalStrength([2, 2]) == 16  # equal elements
assert sol.totalStrength([1, 2, 3]) == 33  # increasing array
assert sol.totalStrength([3, 2, 1]) == 33  # decreasing array

assert sol.totalStrength([5, 5, 5]) == 250  # repeated values

assert sol.totalStrength([1000000000]) == 49  # modulo handling

assert sol.totalStrength([1, 1, 1, 1]) == 20  # duplicate minima

assert sol.totalStrength([9, 1, 7]) == 100  # middle minimum

assert sol.totalStrength([10, 4, 2, 5]) == 307  # mixed ordering
Test Why
[1] Minimum valid input
[1,3,1,2] Official example 1
[5,4,6] Official example 2
[2,2] Duplicate values
[1,2,3] Strictly increasing
[3,2,1] Strictly decreasing
[5,5,5] All equal values
[1000000000] Large number and modulo behavior
[1,1,1,1] Repeated minima
[9,1,7] Middle element dominates
[10,4,2,5] Mixed local minima

Edge Cases

Single Element Array

When the input contains only one wizard, the only subarray is the wizard itself.

For example:

[7]

The answer becomes:

7 × 7 = 49

This case can expose off-by-one mistakes in prefix indexing or boundary initialization. Our implementation handles it correctly because the boundaries default to -1 and n.

Duplicate Values

Equal values are especially dangerous.

For example:

[2,2,2]

Without careful boundary handling, multiple indices may claim the same subarray as their minimum, leading to double counting.

The implementation avoids this by using:

left: strictly smaller
right: smaller or equal

This guarantees every subarray is assigned to exactly one minimum index.

Strictly Increasing or Decreasing Arrays

Arrays like:

[1,2,3,4]

or

[4,3,2,1]

stress monotonic stack logic.

In increasing arrays, most right boundaries extend far.

In decreasing arrays, most left boundaries collapse quickly.

Incorrect stack conditions often fail here. Since each comparison direction is chosen deliberately, the implementation handles both patterns correctly.

Large Values and Overflow

Values can reach:

10^9

and the number of subarrays is enormous.

Intermediate calculations become very large.

The Python version naturally handles big integers, while the Go version explicitly uses int64 and applies modulo arithmetic throughout to prevent overflow and preserve correctness.