利用单调栈的特性
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