Leetcode 42: Trapping Rain Water (2 pass)

leetcode_42b.cpp

#include <iostream>
#include <vector>

using namespace std;

int trap(vector<int>& height) {
    int total = 0, area=0, n=height.size();
    int prevMax=-1, max=-1, pMaxIdx=-1, maxIdx=-1;
    //search forwards
    for (int i=0; i<n; ++i) {
        if (height[i] > 0 && height[i] >= max) {
            prevMax = max;
            pMaxIdx = maxIdx;
            max = height[i];
            maxIdx = i;
            //calc area
            area = 0;
            if (prevMax <= 0 || maxIdx-pMaxIdx <= 1) continue;
            for (int j=pMaxIdx+1; j<maxIdx; ++j) area += (prevMax-height[j]);
            total += area;
        }
    }
    int fwdMax = max;
    prevMax=-1, max=-1, pMaxIdx=-1, maxIdx=-1;
    //search backwards
    for (int i=n-1; i>=0; --i) {
        if (height[i] > 0 && height[i] >= max) {
            prevMax = max;
            pMaxIdx = maxIdx;
            max = height[i];
            maxIdx = i;
            //calc area
            area = 0;
            if (prevMax <= 0 || pMaxIdx-maxIdx <= 1) continue;
            if (prevMax == max && max == fwdMax) continue;
            for (int j=pMaxIdx-1; j>maxIdx; --j) area += (prevMax-height[j]);
            total += area;
        }
    }
    return total;
}

int main() {
    vector<int> heights {2,0,2}; //{0,1,0,2,1,0,1,3,2,1,2,1};
    cout << trap(heights) << "\n";
    return 0;
}