42. Trapping Rain Water

"""
42. Trapping Rain Water
Hard
2489
47
Favorite
Share
Given n non-negative integers representing an elevation map where the width of each bar is 1, compute how much water it is able to trap after raining.

 

The above elevation map 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. Thanks Marcos for contributing this image!

Example:

Input: [0,1,0,2,1,0,1,3,2,1,2,1]
Output: 6
"""

把图形分成两份,前半边是递增,水高为之前的所有数的最大值,后半段是递减,水深是之后所有数的最大值

先写的是用字典来记录水深信息

class Solution:
    def trap(self, height):
        """
        :type height: List[int]
        :rtype: int
        """
        if not height:
            return 0
        maxindex = height.index(max(height))
        dict_height = {-1:0, len(height):0}
        for i in range(maxindex+1):
            dict_height[i] = max(dict_height[i-1], height[i])
        for i in range(len(height)-1, maxindex, -1):
            dict_height[i] = max(dict_height[i+1], height[i])
        sum = 0
        for i in range(len(height)+1):
            sum += dict_height[i] - height[i]
        return sum

 

后来想一想觉得可以改成空间复杂度为O(1)的代码

class Solution:
    def trap(self, height):
        """
        :type height: List[int]
        :rtype: int
        """
        height_len = len(height)
        if height_len < 3:
            return 0
        maxindex = height.index(max(height))
        sum = 0
        tmp = height[0]
        for i in range(1, maxindex):
            tmp = max(tmp, height[i])
            sum += tmp - height[i]
        tmp = height[height_len-1]
        for i in range(height_len-2, maxindex, -1):
            tmp = max(tmp, height[i])
            sum += tmp - height[i]
        return sum