Problem Statement in English
You’re given an array of integers nums, indicating the height of the walls. Your task is to return the total water that can be trapped in between all the walls put together.
Approach
Full disclosure: this is not my code, I got this online from NeetCode (brilliant channel by the way, definitely check him out).
There’s 2 approaches to this, one with $O(n)$ space complexity, and another with $O(1)$ space complexity.
Linear Space Solution: This uses a Prefix Sums based approach. We make a left pass and calculate the heightest wall up until the current index we’re processing, and a right pass which does the same thing, but from the right.
Then for any given index we know the highest wall on its left and right. Using this we can calculate the amount of water that can be held at that index.
Constant Space Solution: This solution uses the 2 Pointer Technique. Check out this post as this question is a slight variation of that one.
The objective is the same as in the linear space solution. We want to always have the highest walls on the left and right of any index that we’re processing. Just that here we don’t store these heights in arrays, hence saving on space.
To decide which pointer to move, we compare the heights at the left and right pointers. We move the pointer with the smaller height, since that is the bottleneck for how much water can be trapped at that index. Once we move the pointer, we calculate the amount of water that can be trapped at that index, add it to our result, and update the left and right max heights accordingly.
Take a look at the code now.
Solution in Python
O(n) space
# O(n) space
class Solution:
def trap(self, height: List[int]) -> int:
N = len(height)
prefix = [0] * N
prefix[0] = height[0]
suffix = [0] * N
suffix[-1] = height[-1]
for i in range(1, N - 1):
prefix[i] = max(prefix[i - 1], height[i])
for i in reversed(range(1, N - 1)):
suffix[i] = max(suffix[i + 1], height[i])
res = 0
for i in range(1, N - 1):
limit = min(prefix[i - 1], suffix[i + 1])
store = limit - height[i]
if store > 0:
res += store
return res
O(1) space
class Solution:
def trap(self, height: List[int]) -> int:
N = len(height)
l, r = 0, N - 1
leftMax = height[l]
rightMax = height[r]
res = 0
while l < r:
if height[l] < height[r]:
l += 1
temp = height[l]
else:
r -= 1
temp = height[r]
depth = min(leftMax, rightMax) - temp
if depth > 0:
res += depth
leftMax = max(leftMax, height[l])
rightMax = max(rightMax, height[r])
return res
Complexity
Time: $O(n)$
Since we process $n$ elements at most.Space: $O(1)$
Since we only use a few integer variables.
Mistakes I Made
I had to look this one up :(
And we are done.