Trapping Rain Water in Python

The trapping rain water problem is a classic algorithmic challenge where we calculate how much water can be trapped after raining on an elevation map represented by an array of heights. Each element represents the height of a bar with width 1.

In this visualization, we can see 6 units of water trapped between the elevation bars, shown in blue.

Algorithm Explanation

The solution uses a stack-based approach to calculate trapped water. The algorithm works by:

• Using a stack to keep track of indices of bars

• When we find a bar taller than the previous one, water can be trapped

• Calculating the trapped water between the current bar and previous bars

Step-by-Step Process

• Initialize a stack, water counter, and index pointer

• Iterate through each height:

  • If stack is empty or current height is smaller, push index to stack

  • If current height is larger, pop from stack and calculate trapped water

  • Water area = distance × (min_height - popped_height)

• Return total trapped water

Implementation

class Solution:
    def trap(self, height):
        stack = []
        water = 0
        i = 0
        
        while i < len(height):
            # If stack is empty or current height is smaller/equal
            if len(stack) == 0 or height[stack[-1]] >= height[i]:
                stack.append(i)
                i += 1
            else:
                # Pop from stack and calculate trapped water
                x = stack.pop()
                
                if len(stack) != 0:
                    # Calculate trapped water
                    min_height = min(height[stack[-1]], height[i])
                    distance = i - stack[-1] - 1
                    water += distance * (min_height - height[x])
        
        return water

# Test the solution
solution = Solution()
heights = [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]
result = solution.trap(heights)
print(f"Input: {heights}")
print(f"Trapped water: {result} units")

Input: [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]
Trapped water: 6 units

How It Works

The algorithm identifies valleys between taller bars where water can accumulate. When we encounter a bar taller than the previous one, we calculate how much water can be trapped in the valley between them.

Example Walkthrough

For the input [0,1,0,2,1,0,1,3,2,1,2,1]:

• Water trapped at index 2: 1 unit (between heights 1 and 2)

• Water trapped at indices 4-6: 3 units (between heights 2 and 3)

• Water trapped at index 9: 1 unit (between heights 2 and 2)

• Total: 1 + 3 + 1 + 1 = 6 units

Time and Space Complexity

• Time Complexity: O(n) - each element is pushed and popped at most once

• Space Complexity: O(n) - stack can store all indices in worst case

Conclusion

The stack-based approach efficiently solves the trapping rain water problem by identifying valleys where water accumulates. This method processes each bar exactly once, making it optimal for large datasets.