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.

Practice

1. What is the main purpose of using a stack in the Trapping Rain Water Using Stack algorithm?

easy

A. To store the total amount of trapped water

B. To sort the heights of bars in ascending order

C. To keep track of bars that are not yet processed for water trapping

D. To reverse the order of the height array

Solution

1. Step 1: Understand the role of the stack

   The stack stores indices of bars that are waiting to find a right boundary to trap water.

2. Step 2: Why not other options?

   The stack does not store water amount or sort bars; it helps find boundaries for trapped water.

3. Final Answer:

   To keep track of bars that are not yet processed for water trapping -> Option C

4. Quick Check:

   Stack tracks bars for trapping water = B [OK]

Hint: Stack holds bars waiting for right boundary [OK]

Common Mistakes:

• Thinking stack stores water amount
• Assuming stack sorts the bars
• Believing stack reverses the array

2. Which of the following Python code snippets correctly initializes an empty stack for the trapping rain water algorithm?

easy

A. stack = 0

B. stack = {}

C. stack = None

D. stack = []

Solution

1. Step 1: Identify correct stack initialization

   In Python, a stack is commonly represented by a list, so stack = [] initializes an empty stack.

2. Step 2: Why other options are wrong

   {} is a dictionary, None is no value, and 0 is an integer, none suitable for stack operations.

3. Final Answer:

   stack = [] -> Option D

4. Quick Check:

   Empty list for stack = A [OK]

Hint: Use empty list [] to create stack in Python [OK]

Common Mistakes:

• Using {} which creates a dictionary
• Assigning None instead of list
• Using integer 0 as stack

3. What is the output of the following code snippet that uses stack to calculate trapped rain water?

def trap(height):
    stack = []
    water = 0
    for i, h in enumerate(height):
        while stack and height[stack[-1]] < h:
            top = stack.pop()
            if not stack:
                break
            distance = i - stack[-1] - 1
            bounded_height = min(height[stack[-1]], h) - height[top]
            water += distance * bounded_height
        stack.append(i)
    return water

print(trap([0,1,0,2,1,0,1,3,2,1,2,1]))

medium

A. 6

B. 5

C. 4

D. 7

Solution

1. Step 1: Trace the algorithm with given heights

   The code calculates trapped water for heights [0,1,0,2,1,0,1,3,2,1,2,1]. The total trapped water is 6 units.

2. Step 2: Confirm calculation matches known example

   This is a classic example where trapped water is 6, matching the output of the code.

3. Final Answer:

   6 -> Option A

4. Quick Check:

   Trapped water = 6 [OK]

Hint: Classic example trap water = 6 units [OK]

Common Mistakes:

• Miscounting trapped water units
• Ignoring stack pop conditions
• Wrong distance or bounded height calculation

4. Identify the error in the following code snippet for trapping rain water using stack:

def trap(height):
    stack = []
    water = 0
    for i, h in enumerate(height):
        while stack and height[stack[-1]] <= h:
            top = stack.pop()
            if not stack:
                break
            distance = i - stack[-1] - 1
            bounded_height = min(height[stack[-1]], h) - height[top]
            water += distance * bounded_height
        stack.append(i)
    return water

print(trap([4,2,0,3,2,5]))

medium

A. The distance calculation is incorrect

B. The condition should be < instead of <=

C. The stack should be a queue

D. The water variable should be initialized to 1

Solution

1. Step 1: Check the while loop condition

   The condition height[stack[-1]] <= h causes popping even when heights are equal, which is incorrect for trapping water calculation.

2. Step 2: Correct condition and reason

   It should be height[stack[-1]] < h to pop only when current height is strictly greater, ensuring correct water trapping.

3. Final Answer:

   The condition should be < instead of <= -> Option B

4. Quick Check:

   Use < not <= in while condition [OK]

Hint: Use strict less than < in stack pop condition [OK]

Common Mistakes:

• Using <= causes wrong popping
• Confusing stack with queue
• Incorrect distance formula
• Initializing water to 1 instead of 0

5. Given the height array [3, 0, 2, 0, 4], what is the total trapped rain water calculated using the stack method? Explain how the stack helps in this calculation.

hard

A. 4

B. 7

C. 5

D. 6

Solution

1. Step 1: Calculate trapped water step-by-step

   Heights: [3,0,2,0,4]. Water trapped at position 1 is 3 units, at position 2 is 0 units, at position 3 is 1 unit. Total trapped water = 4 units.

2. Step 2: Explain stack role in calculation

   The stack stores indices of bars. When a higher bar is found, it pops lower bars and calculates trapped water using distance and bounded height between bars.

3. Final Answer:

   4 -> Option A

4. Quick Check:

   Stack helps find boundaries; trapped water = 4 [OK]

Hint: Stack finds left/right boundaries to calculate trapped water [OK]

Common Mistakes:

• Adding trapped water incorrectly
• Ignoring zero height bars
• Not using stack to find boundaries