Code : Imbalanced Array

#include <bits/stdc++.h>
using namespace std;

// To handle duplicate maximas, assume that the leftmost occurence of the
// maxima is the largest element, and all other occurence are smaller than it.

long long solve(vector<long long> &a) {
    int n = a.size();

    stack<int> st;
    vector<int> nge(n), pge(n);
    for (int i = 0; i < n; i++) {
        while (!st.empty() && a[st.top()] < a[i]) {
            nge[st.top()] = i;
            st.pop();
        }
        if (st.empty()) {
            pge[i] = -1;
        } else {
            pge[i] = st.top();
        }
        st.push(i);
    }
    while (!st.empty()) {
        nge[st.top()] = n;
        st.pop();
    }

    vector<int> nse(n), pse(n);
    for (int i = 0; i < n; i++) {
        while (!st.empty() && a[st.top()] > a[i]) {
            nse[st.top()] = i;
            st.pop();
        }
        if (st.empty()) {
            pse[i] = -1;
        } else {
            pse[i] = st.top();
        }
        st.push(i);
    }
    while (!st.empty()) {
        nse[st.top()] = n;
        st.pop();
    }

    long long ans = 0;
    for (int i = 0; i < n; i++) {
        ans += a[i] * (i - pge[i]) * (nge[i] - i);
        ans -= a[i] * (i - pse[i]) * (nse[i] - i);
    }

    return ans;
}

int main() {
    int n;
    cin >> n;
    vector<long long> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }
    cout << solve(a) << endl;
}