trap.py

class Solution:
    def trap(self, height: List[int]) -> int:
        stack = []
        ans = 0
        for i in range(len(height)):
            if not stack or height[i] < height[stack[-1]]:
                stack.append(i)
            else:
                while stack and height[i] > height[stack[-1]]:
                    cur = stack.pop()
                    if not stack:
                        break
                    h = min(height[i], height[stack[-1]]) - height[cur]
                    w = i - stack[-1] - 1
                    ans += h * w
                stack.append(i)
        return ans