Cow Frisbee (C++11)

Falin Yu (lisayu999)
Submitted: Mon, Jan 31, 2022 00:52:33 EST
Contest: USACO 2022 January Contest, Silver

/*
Jan. 2022, silver, problem 2
*/
#include <bits/stdc++.h>
using namespace std;
#define PP pair<int,int>
#define LL long long
// use stack for left greater and right greater
void xpop(stack<PP>& s,vector<int>& v) {
    PP a=s.top();
    // cout << a.first << " " << a.second << endl;
    s.pop();
    if(s.empty()) {
        // cout << "empty" << endl;
        v[a.second] = -1;
    } else {
        // cout << a.second << endl;
        v[a.second] = s.top().second;
    }
}
int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    int N; 
    stack<PP> sp;
    cin >> N;
    vector<LL> heights(N);
    vector<int> left(N,-1);
    vector<int> right(N,-1);
    for(int i=0;i<N;i++) {
        cin >> heights[i];
    }
    // find left near greater
    for(int i=0;i<N;i++) {
        while(!sp.empty() && sp.top().first < heights[i]) {
            xpop(sp, left);
        }
        sp.push({heights[i],i});
    }
    while(!sp.empty()) {
        xpop(sp, left);
    }
    // find right near greater
    for(int i=N-1; i>=0; i--){
        while(!sp.empty() && sp.top().first < heights[i]) {
            xpop(sp, right);
        }
        sp.push({heights[i],i});
    }
    while(!sp.empty()) {
        xpop(sp, right);
    }
    
    LL total_distance = 0;
    for(int i=0;i<N;i++) {
        // each adjacent contribute 2 distances
        if(i!=0) total_distance += 2; 
        // left greater and right greater forming a parabola
        if( (left[i] != -1)&&(right[i] !=-1)) {
            total_distance = total_distance + (right[i]-left[i]+1);
        }
    }
    cout << total_distance << endl;
}