Skip to content

Trapping rain water

Tim_Gao edited this page Sep 16, 2016 · 1 revision
public class Solution {
    public int trap(int[] height) {
        if(height.length < 3){
            return 0;
        }
        int water = 0;
        //store index in height
        Stack<Integer> st = new Stack<>();
        for (int i=0; i<height.length; ++i){
            if (st.isEmpty()){
                st.push(i);
            } else {
                if (height[i] <= height[st.peek()] ){
                    water += height[i] * (i-st.peek()-1);
                    st.push(i);
                } else {
                    //height[i] > height[st.peek()]
                    int prevHeight = 0;
                    while (!st.isEmpty() && height[i] >= height[st.peek()]){
                        int topIndex = st.pop();
                        water += (height[topIndex] - prevHeight)*(i-topIndex-1);
                        prevHeight = height[topIndex];
                    }
                    if (!st.isEmpty()){
                        water += (height[i]-prevHeight)*(i-st.peek()-1);
                    }
                    st.push(i);
                }
            }
        }
        return water;
    }
}

Clone this wiki locally