Hard

Sum of Total Strength of Wizards

LAB

For every contiguous group, strength is min(group) times sum(group). Return total strength modulo 1000000007.

EXAMPLES

Example 1

Input
{
  "strength": [
    1,
    3,
    1,
    2
  ]
}

Output
44

FUNCTION SHAPE

strength: intArray→int

SOLUTION NOTE

Extension of Q907. Need prefix sum of prefix sums to efficiently compute sum of all subarray sums. Uses PL + NLE for contribution counting.

Reveal reference solution +

pythonREFERENCE

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

    # NLE on the right
    right = [n] * n
    stack = []
    for i in range(n):
        while stack and A[stack[-1]] >= A[i]:
            right[stack.pop()] = i
        stack.append(i)

    # PL on the left
    left = [-1] * n
    stack = []
    for i in range(n-1, -1, -1):
        while stack and A[stack[-1]] > A[i]:
            left[stack.pop()] = i
        stack.append(i)

    # Prefix sum of prefix sums for efficient range sum queries
    prefix = list(accumulate(accumulate(A), initial=0))

    res = 0
    for i in range(n):
        l, r = left[i], right[i]
        # Sum of all subarrays where A[i] is minimum
        # Left choices: (i - l), Right choices: (r - i)
        left_sum = prefix[i + 1] * (r - i) - prefix[r + 1] * (i - l)
        right_sum = prefix[i] * (r - i) - prefix[l] * (r - i)
        contribution = (left_sum - right_sum) * A[i]
        res = (res + contribution) % MOD

    return res % MOD

TimeO(n)

SpaceO(n)

Open on LeetCode ↗