Looking for good programming challenges?

Use the search below to find our solutions for selected questions!

Trapping rain water challenge

Sharing is caring!

Problem statement
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.

Sample input


Sample output


We can solve this problem in \mathcal{O}(n) time by iterating over the array twice. Once from left to right and once from right to left.

Starting at wall i we iterate from left to right until we find a wall j whose height\left[j\right] is equal to or greater than the current height\left[i\right] keeping track of the sum of difference between all smaller walls and wall i in between i and j. If we have stumbled upon a wall j with height\left[j\right] \ge height\left[i\right] we add the collected sum to the overall capacity. We then continue the same procedure starting from wall j. If we reach the end of the array we do a similar procedure iterating from index length(height) - 1 till the 0, that is, backwards. The reason we need to iterate backwards is to collect the water we have missed in case a wall i was followed only by smaller walls (see illustration above with wall 1). Also notice that in this iteration we need to ignore all the walls j with height\left[j\right] = height\left[i\right] because we have already included them in the first iteration.

Full code

class Solution(object):
    def trap(self, height):
        :type height: List[int]
        :rtype: int
        ans = 0

        # Move from left to right
        i = 0
        while i < len(height):
            cap = 0
            l = height[i]
            j = i + 1
            while j < len(height) and height[j] < l:
                cap += (l - height[j])
                j += 1
            if j < len(height) and height[j] >= l:
                ans += cap
            i = j

        # Move from right to left
        i = len(height) - 1
        while i >= 0:
            cap = 0
            r = height[i]
            j = i - 1
            while j >= 0 and height[j] <= r:
                cap += (r - height[j])
                j -= 1
            if j >= 0 and height[j] >= r:
                ans += cap
            i = j

        return ans

def main():
    solution = Solution()
    print(solution.trap([0,1,0,2,1,0,1,3,2,1,2,1])) # 6
    print(solution.trap([4,2,3])) # 1
    print(solution.trap([4,2,0,3,2,5])) # 9
    print(solution.trap([6,8,5,0,0,6,5])) # 13
    print(solution.trap([2,0,2])) # 2

if __name__ == "__main__": main()