Program to find how many total amount of rain we can catch in Python

The rainwater trapping problem involves calculating how much water can be trapped between elevation bars after it rains. Given an array representing bar heights, we need to find the total trapped water volume.

Algorithm Approach

We use a stack-based approach to identify water-trapping positions. The algorithm processes each bar and calculates trapped water when it finds a "valley" pattern ?

Step-by-Step Process

• Initialize a stack to store indices and a water counter
• For each bar, if it's shorter than the previous bar, push its index
• When we find a taller bar, pop from stack and calculate trapped water
• Water area = distance × (minimum height - valley height)

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, push index
            if len(stack) == 0 or height[stack[-1]] >= height[i]:
                stack.append(i)
                i += 1
            else:
                # Found a taller bar, calculate trapped water
                x = stack.pop()
                
                if len(stack) != 0:
                    # Calculate water trapped above height[x]
                    temp = min(height[stack[-1]], height[i])
                    dist = i - stack[-1] - 1
                    water += dist * (temp - height[x])
        
        return water

# Test the solution
solution = Solution()
heights = [2, 5, 2, 0, 5, 8, 8]
result = solution.trap(heights)
print(f"Heights: {heights}")
print(f"Total trapped water: {result}")

Heights: [2, 5, 2, 0, 5, 8, 8]
Total trapped water: 8

How It Works

The algorithm identifies water-trapping areas by finding indices where water can be contained between higher bars on both sides ?

def trap_with_explanation(height):
    stack = []
    water = 0
    i = 0
    
    print(f"Processing heights: {height}")
    
    while i < len(height):
        if len(stack) == 0 or height[stack[-1]] >= height[i]:
            print(f"Push index {i} (height {height[i]}) to stack")
            stack.append(i)
            i += 1
        else:
            x = stack.pop()
            print(f"Pop index {x} (height {height[x]}) from stack")
            
            if len(stack) != 0:
                left = stack[-1]
                temp = min(height[left], height[i])
                dist = i - left - 1
                trapped = dist * (temp - height[x])
                water += trapped
                print(f"Water trapped: {trapped} units (distance: {dist}, height: {temp - height[x]})")
    
    return water

# Example with step-by-step explanation
result = trap_with_explanation([3, 0, 2, 0, 4])
print(f"Total water trapped: {result}")

Processing heights: [3, 0, 2, 0, 4]
Push index 0 (height 3) to stack
Push index 1 (height 0) to stack
Pop index 1 (height 0) from stack
Water trapped: 2 units (distance: 1, height: 2)
Push index 2 (height 2) to stack
Push index 3 (height 0) to stack
Pop index 3 (height 0) from stack
Water trapped: 2 units (distance: 1, height: 2)
Pop index 2 (height 2) from stack
Water trapped: 1 units (distance: 1, height: 1)
Total water trapped: 5

Time and Space Complexity

Conclusion

The stack-based approach efficiently solves the rainwater trapping problem by identifying water-holding valleys between elevation bars. This method processes each element once, making it optimal for large datasets.