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

[6,8,5,0,0,6,5]

Sample output

13

Solution
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()