Suppose we have an array of n non-negative integers. These are representing an elevation map where the width of each bar is 1, we have to compute how much water it is able to trap after raining. So the map will be like −
Here we can see there are 6 blue boxes, so the output will be 6.
To solve this, we will follow these steps −
Define a stack st, water := 0 and i := 0
while i < size of height
if is stack is empty or height[stack top] >= height[i], then push i into stack, increase i by 1
otherwise
x := stack top element, delete top from stack
if stack is not empty, then
temp := min of height[stack top element] and height[i]
dest := i – stack top element – 1
water := water + dist * (temp – height[x])
return water
Example
Let us see the following implementation to get a better understanding −
class Solution(object): def trap(self, height): stack = [] water = 0 i=0 while i<len(height): if len(stack) == 0 or height[stack[-1]]>=height[i]: stack.append(i) i+=1 else: x = stack[-1] stack.pop() if len(stack) != 0: temp = min(height[stack[-1]],height[i]) dist = i - stack[-1]-1 water += dist*(temp - height[x]) return water ob = Solution() print(ob.trap([0,1,0,2,1,0,1,3,2,1,2,1]))
Input
[0,1,0,2,1,0,1,3,2,1,2,1]
Output
6