Given n
non-negative integers representing an elevation map where the width of each bar is 1
, compute how much water it can trap after raining.
Example 1:
Input: height = [0,1,0,2,1,0,1,3,2,1,2,1] Output: 6 Explanation: The above elevation map (black section) is represented by array [0,1,0,2,1,0,1,3,2,1,2,1]. In this case, 6 units of rain water (blue section) are being trapped.
Example 2:
Input: height = [4,2,0,3,2,5] Output: 9
Constraints:
n == height.length
1 <= n <= 2 * 104
0 <= height[i] <= 105
题目给定一个非负整数的数组,数组里每个数表示一个宽度为1的柱子的高度,由于柱子之间的高度差就会形成一个个槽位,这样一下雨这些槽位就会装满水。问最多能装多少单位的水?
先来分析一下什么情况才能装水,两根紧挨着的柱子肯定无法形成槽位,因此至少是三根以上的柱子;要是柱子高度单调递增或递减也无法形成槽位。只有出现一排柱子边上的两根柱子高于中间的柱子才能形成槽位。根据以上分析,跟单调栈的算法很相似,因此可以用单调栈来解答该题。
遍历数组,只要柱子高度是递减的(无法形成有效槽位)就直接压入栈(压入栈的是数组索引值,因为计算水量是高度乘以宽度,索引值方便计算柱子间的距离),直到碰到当前柱子比前一个(栈顶)柱子高的时候,这种情况要是前一根柱子前还有前一根柱子(堆栈里有两个以上的数),由于是单调递减栈,那就形成了水槽。假设当前柱子位置为i高度为h, 栈顶两根柱子位置和高度分别为i1, i2和h1, h2,水槽的蓄水量为(min(h, h2) - h1) * (i - i2 - 1),计算完水量后栈顶的柱子(i1, h1)就可以被弹出栈,相当于该水槽被土填上了,新的最低高度为(i2, h2)。继续用相同方法把当前柱子跟新的栈顶柱子比较并计算水量,直到当前柱子比栈顶的柱子低时把当前柱子压入栈。整个数组遍历完之后就可以得到总的需水量。
注:当前柱子高度跟栈顶柱子相等时也可以按以上方法处理,相当于是一种特殊情况高度差为0,需水量为0。
class Solution:
def trap(self, height: List[int]) -> int:
n = len(height)
res = 0
st = []
for i in range(n):
while st and height[st[-1]] <= height[i]:
j = st.pop()
if st:
res += (min(height[st[-1]], height[i]) - height[j]) * (i - st[-1] - 1)
st.append(i)
return res