Problem Description

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

Example 1

Input: height = [0,1,0,2,1,0,1,3,2,1,2,1]
Output: 6
Explanation: The above elevation map (black section) 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.

Example 2

Input: height = [4,2,0,3,2,5]
Output: 9

Solution

Before going on to think about an algorithm, we have to first fully understand a sub-problem: how much water can for example the position i hold? It depends on the max height on its left and right sides globally. Now, the question becomes to find left max and right max heights for each positions, which can be achieved by using two arrays to record the left max and right max accordingly.

Based on the above idea, we can even optimize the space complexity to O(1). Since the amount of water the position i can hold depends on the smaller one of the left max and right max: say given the index i, we know the left max before the i, and we also know there is a height on the right that is larger than this left max, now we are ready to compute the amount of water at i. Note that this larger right height does NOT need to be the right max, as long as larger than the left max. So we get the algorithm: we initialize two pointers l and r at the 0 and len-1, maintain two variables left_max and right_max, compare two heights at these two pointers l and r, always move the pointer pointing to the smaller height, at the same time update left_max and right_max. Note that left_max is the left max height at l, while right_max is the right max height at r. This way, at each step, we can compute the amount of water at the whichever position has the less height at l or r. Why? Say if height[l] < height[r], we then also know height[r] > left_max, since if left_max at j(j<l) is larger, then we would have keep moving r until height[r]>height[j] when we are at j. Therefore, we can always compute the amount of water at the whichever position has the less height at l or r.

The third thought is based on this: in order to have the position i hold water, there should be left height and right height that are both larger than the height[i]. We can use a stack to maintain the left height: meaning we keep pushing heights if they are smaller than the top element of stack, that is, the heights are in a descending order while we push it. If we have a height height[i] larger than the top element (smallest) of the stack, the height height[i] could be potentially a right height. Since the top element is always the smallest height in the stack, we know it must have a larger left height (assume the stack has more than 1 elements), thus the height at i and the second smallest element in the stack will be the right and left height correspondingly. The amount of water it hold depends on the smaller height of these two heights. And then we remove the top element, keep comparing the next smallest element in the stack to height[i], and compute the amount of water if it is still smaller than height[i], until we meet some element in the stack which is at least no smaller than height[i]. As we can see, heights in the stack are in a non-ascending order. In the end, if we put it graphically, the way we are computing the water is to compute layer by layer from the lower to higher layers.

Java Code

Two Arrays to keep left max and right max

class Solution {
    public int trap(int[] height) {
        int len = height.length;
        int[] left_max = new int[len], right_max = new int[len];
        left_max[0] = height[0];
        right_max[len-1] = height[len-1];
        for(int i=1; i<len; i++){
            left_max[i] = Math.max( left_max[i-1], height[i]);
            right_max[len-1-i] = Math.max( right_max[len-i], height[len-1-i]);
        }
        int water = 0;
        for(int i=0; i<len; i++){
            water += Math.max(0, Math.min(left_max[i], right_max[i]) - height[i]);
        }
        return water;
    }
}

Two pointers to keep left max and right max

class Solution {
    public int trap(int[] height) {
        int len = height.length, left = 0, right = len-1;
        int water = 0, left_max = height[0], right_max = height[len-1];
        while(left<right){
            if(height[left] < height[right]){
                if(left_max < height[left]) left_max = height[left];
                water += Math.max(0, left_max - height[left]);
                left ++;
            }else{
                if(right_max < height[right]) right_max = height[right];
                water += Math.max(0, right_max - height[right]);
                right --;
            }
        }
        return water;
    }
}

Stack to compute layer by layer

class Solution {
    public int trap(int[] height) {
        int len = height.length, water = 0, i=1;
        Stack<Integer> stk = new Stack<Integer>();
        stk.push(0);
        while( i<len ){
            while( height[i] > height[stk.peek()] ){ //have larger right bound
                int tmp = stk.pop();
                if( stk.isEmpty() ) break; // there is no left bound
                water += (i-stk.peek()-1) * (Math.min(height[i], height[stk.peek()]) - height[tmp]);
            }
            stk.push( i++ );
        }
        return water;
    }
}