LeetCode: Trapping Rain Water

Given n non-negative integers representing an elevation map where the width of each bar is 1, compute how much water it is able to trap after raining.

For example,
Given [0,1,0,2,1,0,1,3,2,1,2,1], return 6.

The above elevation map is represented by array [0,1,0,2,1,0,1,3,2,1,2,1]. In this case, 6 units of rain water (blue section) are being trapped. Thanks Marcos for contributing this image!

---

We can find the left bound and right bound for each slot when iterating the array.

C++ Code:

/*
 * func: trapping_rain_water
 * goal: Given n non-negative integers representing an elevation map 
 *       where the width of each bar is 1, compute how much water it 
 *       is able to trap after raining.
 * @param A: input array A
 * @param n: size of the array
 * return: the maximum volume
 */
/*
 * Set up left bound and right bound and keep finding available container
 * complexity: time O(n), space O(1)
 */
int trap(int A[], int n){
    assert(n>=0);
    for(int i=0; i<n; i++){
        assert(A[i] >= 0);
    }
    
    //Initialize left bound and right bound
    int left = 0;
    int right = n-1;
    int max_volume = 0;
    while(left < right){
        //If A[left] <= A[right], there may be water between them
        if(A[left] <= A[right]){
            //keep finding the next left bound and adding water at the same tiem
            int k = left + 1;
            //If left bound is larger than current bound, we can store water
            while(A[left] > A[k]){
                max_volume += A[left] - A[k];
                ++k;
            }
            left = k;
        }else{
            int k = right-1;
            while(A[right] > A[k]){
                max_volume += A[right] - A[k];
                --k;
            }
            right = k;
        }
    }
    return max_volume;
}

Python Code:

# func: Given n non-negative integers representing an elevation map
#       where the width of each bar is 1, compute how much water it
#       is able to trap after raining.
# @param A: a list of integers
# @return: the maximum value
def trap_water(A):
    assert (x >= 0 for x in A)

    left = 0
    right = len(A)-1
    max_volume = 0
    while left < right:
        if A[left] <= A[right]:
            k = left + 1
            while A[k] < A[left]:
                max_volume += A[left] - A[k]
                k += 1
            left = k
        else:
            k = right - 1
            while A[k] < A[right]:
                max_volume += A[right] - A[k]
                k -= 1
            right = k

    return max_volume