Leetcode42 接雨水 Python

利用单调栈的特性

class Solution:
    """
    单调栈,维护一个从大到小的单调栈，比如[4, 3, 2, 1]
    假设我们的数据是[4, 3, 1, 0, 1, 2, 4]
    我们先组成单调栈,则stack里的元素是[4, 3, 1, 0]
    然后我们遇到1，它大于我们栈顶元素
    此时先pop出栈顶元素，记为cur
    现在栈顶元素是左边的水柱
    右边的柱子是当前元素
    然后用min取个最小值，减去cur的高度，就是存储水量，最后乘上两个水柱的距离
    将元素入栈
    
    stack = [4, 3, 1, 1]
    此时遇到2
    我们有2个1，所以都pop出去
    然后根据上面的流程计算，再入栈
    
    stack = [4, 3, 2]
    以此类推
    
    
    
    """
    def trap(self, height):
        res = 0
        stack = [] # 第一维为数字，第二维是索引
        for i in range(len(height)):
            while stack and stack[-1][0] < height[i]:
                cur = stack.pop()
                while stack and cur[0] == stack[-1][0]:
                    # 如果都是相同的元素则继续pop
                    stack.pop()
                if stack:
                    # 先取左右两边最小的柱子，再减去水柱的高度，就是能存储的雨水量
                    # 最后乘两个距离，记得-1
                    res += ((min(stack[-1][0], height[i]) - cur[0])*(i-stack[-1][1]-1))
            stack.append([height[i], i])
        return res