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]