Framework Thinking - Neetcode 150 - Trapping Rain Water

1. Problem understanding

  • Given an array height[] where each element represents an elevation bar of width 1, compute how much water it can trap after raining.

  • You return a single integer — total units of trapped water.

2. Clarify questions

  • Input size n = len(height). Typical constraints: 0 <= n <= 10^5 (so aim for O(n) time, O(1) extra space).

  • Values: 0 <= height[i] <= 10^9 (so use integers).

  • Edge cases: empty array [] → 0; single element or two elements → 0.

  • Duplicates allowed.

  • Negative numbers? Usually not, but if present treat them as negative heights (same logic applies); we’ll assume non-negative heights.

  • Want minimum memory? If yes, prefer the two-pointer O(1) space solution.

  • Want simplicity over micro-optimizations? Precompute left/right max arrays (O(n) time, O(n) space) is simple and robust

3. Work through examples

  • Example A: [0,1,0,2,1,0,1,3,2,1,2,1] — classic answer 6.

  • Example B: [4,2,0,3,2,5] — answer 9 (water at indexes: 1→2, 2→4, 3→1, 4→2, total 9).

  • Example C: [] → 0.

  • Example D: [2] → 0.

4. Brainstorm 2 - 3 solutions

4.1. Brute force & Ideas

 4|                #
 3|#      ~ ~ ~ ~  #
 2|#  ~ ~ # ~ ~ #  #
 1|#  # # # # # #  #
 0|# # # # # # # # #
   0 1 2 3 4

water[i] += min(leftMax, rightMax) - height[i]

Observation

  • If the left is smaller than height[i] => Can not keep water in the left.

  • if the right is smaller than height[i] => Can not keep water in the right.

Why we use leftMax and rightMax, but not using height[0] and height[2] for the water store at block 1 ?

  • For example, at index 1 in the ABOVE EXAMPLE.

  • At index 1, we can store, (4 - 3) - 0 = 3 water index due to the overflow the water, shrink from the height[4] to height[0]

What actually store at height[1], it is water[1] ?

  • water[i] is the water can trapped in the block i.

  • leftMax and rightMax including index i

In case:

  • i is leftMax < rightMax, water = 0.

  • i is rightMax > leftMax, water = leftMax - height[i] = 0 (due to leftMax < rightMax, but rightMax is height[i])

4.2. Brute Force (Naive)

  • Using two loop, calculate leftMax and rightMax each time in the loop => repeatable than pre-computed leftMax and rightMax array.

  • Time: O(N^2).

  • Space: O(1).

4.3. Prefix & Suffix Arrays (Pre-computed leftMax and rightMax)

  • Pre-computed prefix and suffix array => O(N).

  • Time: O(N).

  • Space: O(N).

4.4. Stack

  • Same idea, but calculate dynamically leftMax > mid, mid < height[i] (do not is rightMax)

  • Current situation is

height[left] > height[mid]
height[right] > height[mid]
 left        right
   |            |
   |    ___     |
   |___|   |____|
     mid
  • Fomula:
h = min(right, left) - mid
w = i - stack[-1] - 1
water += h * w
  • Time: O(N).

  • Space: O(N).

4.5. Two Pointers

  • Using two pointer to find the dynamic min(leftMax, rightMax) - height[i].

  • Time: O(N).

  • Space: O(1).

5. Implement solutions

5.1. Brute Force (Naive solution)

class Solution:
    def trap(self, height: List[int]) -> int:
        if not height:
            return 0
        n = len(height)
        res = 0

        for i in range(n):
            leftMax = rightMax = height[i]

            for j in range(i):
                leftMax = max(leftMax, height[j])
            for j in range(i + 1, n):
                rightMax = max(rightMax, height[j])

            res += min(leftMax, rightMax) - height[i]
        return res
  • Time: O(N^2).

  • Space: O(1).

5.2. Prefix & Suffix Arrays (Pre-computed leftMax and rightMax array)

class Solution:
    def trap(self, height: List[int]) -> int:
        n = len(height)
        if n == 0:
            return 0

        leftMax = [0] * n
        rightMax = [0] * n

        leftMax[0] = height[0]
        for i in range(1, n):
            leftMax[i] = max(leftMax[i - 1], height[i])

        rightMax[n - 1] = height[n - 1]
        for i in range(n - 2, -1, -1):
            rightMax[i] = max(rightMax[i + 1], height[i])

        res = 0
        for i in range(n):
            res += min(leftMax[i], rightMax[i]) - height[i]
        return res
  • Time: O(N).

  • Space: O(N).

5.3. Stack (Another approach with mid < leftMax, mid < height[i])

class Solution:
    def trap(self, height: List[int]) -> int:
        if not height:
            return 0
        stack = []
        res = 0

        for i in range(len(height)):
            while stack and height[i] >= height[stack[-1]]:
                # mid < leftMax
                mid = height[stack.pop()]
                if stack:
                    # mid < height[i] 
                    right = height[i]
                    left = height[stack[-1]]

                    # Dynamic calculate
                    h = min(right, left) - mid
                    w = i - stack[-1] - 1
                    
                    # Sum
                    res += h * w
            stack.append(i)
        return res
  • Time: O(N).

  • Space: O(N).

5.4. Two Pointers (Dynamic Calculate leftMax and rightMax without store)

class Solution:
    def trap(self, height: List[int]) -> int:
        if not height:
            return 0

        l, r = 0, len(height) - 1
        leftMax, rightMax = height[l], height[r]
        res = 0
        while l < r:
            if leftMax < rightMax:
                l += 1
                leftMax = max(leftMax, height[l])
                res += leftMax - height[l]
            else:
                r -= 1
                rightMax = max(rightMax, height[r])
                res += rightMax - height[r]
        return res
  • Time: O(N).

  • Space: O(1).

6. Run the testcase

6.1. Stack

height = [3, 0, 2, 0, 4]
  • i = 0, stack = [0]

  • i = 1, 0 < 3 => stack [0, 1].

  • i = 2, 2 > 0 => stack [0, 2], res = 2 * (3 - 1 -1) = 2.

6.2. Leftmax and rightmax

 4|                #
 3|#      ~ ~ ~ ~  #
 2|#  ~ ~ # ~ ~ #  #
 1|#  # # # # # #  #
 0|# # # # # # # # #
   0 1 2 3 4

water[i] += min(leftMax, rightMax) - height[i]
Last Updated On October 27, 2025