Python实现接雨水问题

It’s a problem you are likely to be asked when you have interview from ByteDance.

描述

给定 n 个非负整数表示每个宽度为 1 的柱子的高度图，计算按此排列的柱子，下雨之后能接多少雨水。

上面是由数组 [0,1,0,2,1,0,1,3,2,1,2,1] 表示的高度图，在这种情况下，可以接 6 个单位的雨水（蓝色部分表示雨水）。

解题思路

用双指针来做，左右分别一个指针，一个最大值。即left，right和left_max，right_max
由于较小的峰值决定可以存多少水，所以较小的那边移动指针。
例如移动左边指针时（最大值也值左边最大值）：
遇到比当前最大值大的值时，更新最大值；
遇到比当前最大值小的值，将结果加入最后结果。

代码例子

class Solution:
    def trap(self, A):
        n = len(A)
        if n<=2:
            return 0
        i, j = 0, n - 1
        res = 0
        left_max=0
        right_max=len(A) - 1
        while i < j:
            if A[i] > A[j]: # 右边低于左边
                if A[right_max]<A[j]: # 右边当前高度小于右边最大高度
                    right_max = j # 高度差为当前柱子上面的水量
                    j -= 1
                else:
                    res += A[right_max] - A[j] # 否则更新右边最大高度
                    j -= 1
            else: #  左边低于右边
                if A[left_max] < A[i]:
                    left_max = i
                    i += 1
                else:
                    res += A[left_max]-A[i]
                    i += 1
        return res