A stack based solution for reference, inspired by Histogram


  • 89
    A

    Indeed this question can be solved in one pass and O(1) space, but it's probably hard to come up with in a short interview. If you have read the stack O(n) solution for Largest Rectangle in Histogram, you will find this solution is very very similar.

    The main idea is : if we want to find out how much water on a bar(bot), we need to find out the left larger bar's index (il), and right larger bar's index(ir), so that the water is (min(A[il],A[ir])-A[bot])*(ir-il-1), use min since only the lower boundary can hold water, and we also need to handle the edge case that there is no il.

    To implement this we use a stack that store the indices with decreasing bar height, once we find a bar who's height is larger, then let the top of the stack be bot, the cur bar is ir, and the previous bar is il.

    public int trap(int[] A) {
            if (A==null) return 0;
            Stack<Integer> s = new Stack<Integer>();
            int i = 0, maxWater = 0, maxBotWater = 0;
            while (i < A.length){
                if (s.isEmpty() || A[i]<=A[s.peek()]){
                    s.push(i++);
                }
                else {
                    int bot = s.pop();
                    maxBotWater = s.isEmpty()? // empty means no il
                    0:(Math.min(A[s.peek()],A[i])-A[bot])*(i-s.peek()-1);
                    maxWater += maxBotWater;
                }
            }
            return maxWater;
        }

  • 1
    L

    I like this solution very well! Please see the C++ version in follows:

    class Solution {
    public:
        int trap(vector<int>& height) {
            stack<int> s;
            int max_water(0);
            int i(0);
            while (i<height.size()) {
                if (s.empty()||height[i]<=height[s.top()]) {
                    s.push(i++);
                } else {
                    int bot=height[s.top()];
                    s.pop();
                    max_water+=s.empty()?0:((min(height[i],height[s.top()])-bot)*(i-s.top()-1));
                }
            }
            return max_water;
        }
    };

  • 0
    F

    What is the relation between the names of your variables and the actual things ? For example why index is "ir" and bar "bot" ?


  • 3
    E

    @federico I suppose the "il" means "index of left" and "bot" means "bottom" which is the bottom of the area calculated. Hope this helps you.


  • 0
    G

    Thank you for sharing! I like your solution and here is the Python version:

    class Solution(object):

    def trap(self, height):
        stack = list()
        index = 0
        water = 0
        while index < len(height):
            if len(stack) == 0 or height[index] <= height[stack[len(stack)-1]]:
                stack.append(index)
                index += 1
            else:
                bottom = stack.pop()
                if len(stack) != 0:
                    water += (min(height[stack[len(stack)-1]], height[index]) - height[bottom]) * (index - stack[len(stack)-1] - 1)
        return water

  • 2
    B

    @Eddie-wang I assume ir is index_right and il is index_left. bot is the bottom of the "water container"


  • 0
    J

    Thanks for sharing the code, it helps me find the mistakes in my solution


  • 0
    B

    Great!!!! Nice solution!


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.