rain-trap

Trapping Rain Water

Brute Force

Iterate over each height, calculate the maximum amount of water it can hold, which is the min of max height to the left and that to the right, and sum them up.

This is a $O(n^2)$ solution.

vol = 0
for i in range(1, len(height) - 1):  # the ends can't hold water anyway
    l_max = max(height[:i + 1])  # the current height included
    r_max = max(height[i:])
    vol += min(l_max, r_max) - height[i]
return vol

DP

Apparently, the l_max and r_max are used repeatedly, and we can save calculation by memoizing them. This algorithm is $O(n)$ in both time and space.

vol = 0
n = len(height)
l_max, r_max = [0] * n, [0] * n
# set the ends
l_max[0] = height[0]
r_max[n - 1] = height[n - 1]
for i in range(1, n):
    l_max[i] = max(height[i], l_max[i - 1])
for i in range(n - 2, -1, -1):
    r_max[i] = max(height[i], r_max[i + 1])
for i in range(1, n - 1):
    vol += min(l_max[i], r_max[i]) - height[i]
return vol

Stack

Maintain a stack to keep track of the current vol. This algorithm is $O(n)$ in both time and space.

vol = i = 0
stack = []
n = len(height)
while i < n:
    while stack and height[i] > height[stack[-1]]:
        # the new height is greater than the one at top,
        # so the top is bounded from right
        top = stack.pop()
        if not stack:
            # the stack is empty -- no water can be caught
            break
        # to address the case of adjacent same heights
        distance = i - stack[-1] - 1
        # height[i]         => from right
        # height[stack[-1]] => from left
        # height[top]       => the bottom
        bounded_height = (
            min(height[i], height[stack[-1]]) - height[top]
        )
        vol += distance * bounded_height
    stack.append(i)
    i += 1
return vol

Two Pointers

This algorithm runs in $O(n)$ and $O(1)$ space. The good thing is we approach from both ends, and alternate between left and right pointers to get the volume of water that we know for sure is caught, and avoided handling edge cases at the end.

# approach from both ends
l, r = 0, len(height) - 1
vol = 0
l_max, r_max = 0, 0
while l < r:
    if height[l] < height[r]:
        # get the vol caught from left
        l_max = max(l_max, height[l])  # calculate l_max on the flight
        vol += l_max - height[l]
        l += 1
    else:
        r_max = max(r_max, height[r])
        vol += r_max - height[r]
        r -= 1
return vol