Trapping Rain Water Using Stack in DSA Python

Mental Model

We use a stack to remember bars that can trap water when a taller bar comes after them. The stack helps find boundaries to hold water between bars.

Analogy: Imagine a row of walls of different heights. When rain falls, water collects in dips between taller walls. The stack helps us find these walls quickly, like marking where water can stay.

Height bars: 0 1 0 2 1 0 1 3 2 1 2 1
Stack stores indices of bars that might trap water
Example:
Index: 0 1 2 3 4 5 6 7 8 9 10 11
Bars:  0 1 0 2 1 0 1 3 2 1  2  1
Stack: [] initially empty

Dry Run Walkthrough

Input: height = [0,1,0,2,1,0,1,3,2,1,2,1]

Goal: Calculate total water trapped between bars after raining

Step 1: Push index 0 (height 0) onto stack

Stack: [0]
Water trapped: 0

Why: Start with first bar as potential boundary

Step 2: i=1 height 1 > height[0]=0, pop 0 (stack empty after pop, no water calc), push 1

Stack: [1]
Water trapped: 0

Why: Current taller than previous, pop it as no trapping possible without left boundary

Step 3: Push index 2 (height 0) since 0 < height[1]=1

Stack: [1,2]
Water trapped: 0

Why: Lower bar, might trap water later

Step 4: i=3 height 2 > height[2]=0, pop 2: left=1 (h=1), dist=3-1-1=1, bounded_h=min(1,2)-0=1, water +=1 (total 1); then 2 > height[1]=1, pop 1: stack empty, break; push 3

Stack: [3]
Water trapped: 1

Why: Right boundary found for popped bars; water only for bar 2, bar 1 has no left boundary

Step 5: Push index 4 (height 1) since 1 < 2

Stack: [3,4]
Water trapped: 1

Why: Lower bar pushed for future trapping

Step 6: Push index 5 (height 0) since 0 < 1

Stack: [3,4,5]
Water trapped: 1

Why: Lower bar pushed for future trapping

Step 7: i=6 height 1 > height[5]=0, pop 5: left=4 (h=1), dist=6-4-1=1, bounded_h=min(1,1)-0=1, water +=1 (total 2); 1 == height[4]=1, no more pop; push 6

Stack: [3,4,6]
Water trapped: 2

Why: Water trapped above bar 5 between 4 and 6

Step 8: i=7 height 3 > height[6]=1, pop 6: left=4 (h=1), dist=7-4-1=2, bounded_h=min(1,3)-1=0, no water; > height[4]=1, pop 4: left=3 (h=2), dist=7-3-1=3, bounded_h=min(2,3)-1=1, water +=3 (total 5); > height[3]=2, pop 3: stack empty, break; push 7

Stack: [7]
Water trapped: 5

Why: Multiple pops: no water above 6 (bounded by 1), water above 4 (and previous uncovered) bounded by 3 and 7

Step 9: Push index 8 (height 2) since 2 < 3

Stack: [7,8]
Water trapped: 5

Why: Lower bar pushed

Step 10: Push index 9 (height 1) since 1 < 2

Stack: [7,8,9]
Water trapped: 5

Why: Lower bar pushed

Step 11: i=10 height 2 > height[9]=1, pop 9: left=8 (h=2), dist=10-8-1=1, bounded_h=min(2,2)-1=1, water +=1 (total 6); 2 == height[8]=2 no more; push 10

Stack: [7,8,10]
Water trapped: 6

Why: Water trapped above bar 9

Step 12: Push index 11 (height 1) since 1 < 2

Stack: [7,8,10,11]
Water trapped: 6

Why: End of array, lower bar pushed

Result:

Final stack: [7,8,10,11]
Total water trapped: 6 units

Annotated Code

DSA Python

class Solution:
    def trap(self, height: list[int]) -> int:
        stack = []  # stores indices
        water_trapped = 0
        for i, h in enumerate(height):
            # While current bar is taller than top of stack bar
            while stack and height[stack[-1]] < h:
                top = stack.pop()  # bar that can trap water
                if not stack:
                    break  # no left boundary
                distance = i - stack[-1] - 1  # width between boundaries
                bounded_height = min(height[stack[-1]], h) - height[top]
                water_trapped += distance * bounded_height
            stack.append(i)
        return water_trapped

# Driver code
height = [0,1,0,2,1,0,1,3,2,1,2,1]
sol = Solution()
print(sol.trap(height))

while stack and height[stack[-1]] < h:

Check if current bar is taller than top stack bar to trap water

top = stack.pop()

Pop the bar that can trap water between boundaries

if not stack: break

No left boundary means no trapped water here

distance = i - stack[-1] - 1

Calculate width between left and right boundaries

bounded_height = min(height[stack[-1]], h) - height[top]

Calculate height of trapped water bounded by shorter boundary

water_trapped += distance * bounded_height

Add trapped water volume for this segment

stack.append(i)

Push current bar index as potential boundary

Output

6

Complexity Analysis

Time: O(n) because each bar is pushed and popped at most once

Space: O(n) for the stack storing indices of bars

vs Alternative: Compared to brute force O(n^2) checking left and right max for each bar, stack method is much faster

Edge Cases

Empty height list

Returns 0 as no bars to trap water

DSA Python

for i, h in enumerate(height):

All bars same height

No water trapped, returns 0

DSA Python

while stack and height[stack[-1]] < h:

Strictly increasing or decreasing bars

No trapped water, returns 0

DSA Python

while stack and height[stack[-1]] < h:

When to Use This Pattern

When you need to find trapped water between bars and want to avoid checking all pairs, use a stack to track boundaries and calculate trapped water efficiently.

Common Mistakes

Mistake: Not checking if stack is empty before accessing stack[-1]
Fix: Add a condition to break if stack is empty after popping

Mistake: Calculating distance incorrectly by not subtracting 1
Fix: Use distance = i - stack[-1] - 1 to get correct width between boundaries

Mistake: Using max instead of min to find bounded height
Fix: Use min(height[stack[-1]], current height) to find correct bounded height

Summary

Calculates how much rainwater can be trapped between bars using a stack to find boundaries.

Use when you want an efficient way to compute trapped water without checking all pairs of bars.

The key insight is using a stack to remember bars that can trap water when a taller bar appears on the right.