Leetcode 907: Sum of Subarray Minimums

https://leetcode.com/problems/sum-of-subarray-minimums/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.

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

Notes:

Since the size of the array is very large, the naive O(N^2) method cannot pass the OJ.

The optimal solution can be O(N), by using a stack.

One of keys to this question is to find the minimum element and also its range.

Let say there is one local minimum arr[i], and the smaller first element on the left side is at index j, and the first smaller element on the right side is at index k. Then there are (i - j) * (k - i) subarrays with arr[i] as the minimum. So the contribution to the final results is arr[i]*(i-j)*(k-j).

And we need to do this for every element. A stack can be used to find j and k for each arr[i] by one pass.

The key is to maintain a monotonically increasing stack.

Since we need to figure out the distance, the stack will store index instead of values.

When meet with a smaller element (< stack.top()), the first smaller element on the right side of the arr[stack.top()] element can be found directly;

After popping out all the larger ones, the remaining are smaller or equal ones. So the current stack.top() is the left boundary of arr[i]. (The equal situation can also be included correctly.)

In order to easy code, element 0 has been placed before and after the array. In this way, all the elements in the original array can be popped out at the end of the code.

See the code below:

class Solution {
public:
    int sumSubarrayMins(vector<int>& A) {
        int res = 0, mod = 1e9+7;
        A.insert(A.begin(), 0);
        A.push_back(0);
        vector<int> L_idx(A.size(), -1), R_idx(A.size(), -1);
        stack<int> sk;
        for(int i=0; i<A.size(); ++i) {
            while(sk.size() && A[sk.top()] > A[i]) {
                int t = sk.top();
                sk.pop();
                R_idx[t] = i;
                long dis1 = t - L_idx[t], dis2 = R_idx[t] - t, val = A[t];
                res = (res + (dis1 * dis2 * val) % mod) % mod;
            }
            if(sk.size()) L_idx[i] = sk.top();
            sk.push(i);
        }
        A.erase(A.begin());
        A.pop_back();
        return res;
    }
};