907. Sum of Subarray Minimums (Python)

Related Topic

Stack.

Description

Given an array of integers A, find the sum of min(B), where B ranges over every (contiguous) subarray of A.

Since the answer may be large, return the answer modulo 10^9 + 7.

Sample I/O

Example 1

Input: [3,1,2,4]
Output: 17
Explanation: Subarrays are [3], [1], [2], [4], [3,1], [1,2], [2,4], [3,1,2], [1,2,4], [3,1,2,4]. 
Minimums are 3, 1, 2, 4, 1, 1, 2, 1, 1, 1.  Sum is 17.

Note

1. 1 <= A.length <= 30000
2. 1 <= A[i] <= 30000

Methodology

Maintain minimum stack is great way to solve this question. We create a stack to store the minimun value. Iterate the given list, since the answer at least inclued every number so for each iteration we init the count as 1. If stack has no value, then we append the current value and count to the stack and update the total_sum and pre_sum. If stack is not empty then we need to compare current value with the top value of stack. There are 2 conditions here:

1. If the top value is smaller or equal to current value, we just append the current value to stack and update the total_sum and pre_sum and continue to next iteration.
2. If top value is greater than current value, to maintain the minimum stack, we need to pop out the top value and we need minus top value * it’s count from pre_sum and continue checking top value of stack until we find the top is equal or smaller than the current. Then we append the current value to stack and update the total_sum and pre_sum.

def sumSubarrayMins(self, A: List[int]) -> int:
        MOD = 10**9+7
        stack = []
        ans = pre_sum = 0
        for index, value in enumerate(A):
            count = 1
            while stack and stack[-1][0]>=value:
                v, c = stack.pop()
                count+=c
                pre_sum-=v*c
            stack.append((value,count))
            pre_sum+=value*count
            ans+=pre_sum
        return ans % MOD

BigO

Time complexity is O(n) where n is size of A