python-收集雨水问题

问题描述：
收集雨水问题：给定n个非负整数，表示直方图的方柱的高度，同时，每个方柱的宽度假定都为1。若使用这样形状的容器收集雨水，可以盛多少水量？
如输入：0,1,0,2,1,0,1,3,2,1,2,1；返回6。

解题思路：
当水盛满的时候，雨水的面积 = 总面积 - 直方图的面积
从左向右，找比当前柱子A高的第一个柱子B，把柱子A和B之间填充上水，既设置他们之间的柱子高度变为柱子A的高度
然后在向右找比柱子B高的第一个柱子C，把柱子B和C之间填充上水，既设置他们之间的柱子高度变为柱子B的高度……
当两个柱子相邻时不用填充
为了防止漏掉还要从右向左执行一遍
最后雨水的面积 = 总面积 - 直方图的面积
代码实现：

from copy import deepcopy
import matplotlib.pyplot as plt


def get_rain(li):
    newli = deepcopy(li)
    size = len(newli)
    # 从左向右填充雨水
    for i in list(range(size - 1)):
        for j in list(range(i + 1, size)):
            if newli[j] >= newli[i]:
                newli[i + 1:j] = [newli[i]] * (j - i - 1)
                break
    # 从右向左填充雨水
    for i in list(range(size - 1, 0, -1)):
        for j in list(range(i - 1, -1, -1)):
            if newli[j] >= newli[i]:
                newli[j + 1:i] = [newli[i]] * (i - j - 1)
                break
    return newli


if __name__ == '__main__':
    li = [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]
    result = get_rain(li)
    print(result)
    print("雨水的面积: %d" % (sum(result) - sum(li)))

    # 画图
    fig = plt.figure(1)
    ax = fig.add_subplot(111)
    fig.set_size_inches(len(li), max(li))
    left = list(range(len(li)))
    plt.bar(left=left, height=result, width=0.95, color='b', alpha=0.2)
    plt.bar(left=left, height=li, width=0.95, color='g', alpha=1)
    plt.xticks(list(range(0, len(li), 1)))
    plt.yticks(list(range(0, max(li) + 2, 1)))
    plt.show()

运行结果：

参考自https://blog.csdn.net/weixin_42018258/article/details/80626637?depth_1-utm_source=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-1&utm_source=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-1以学习记录。