Problem overview
The function trap receives a slice height of non-negative integers, an elevation map in which every bar is 1 unit wide. After rain, water collects in the dips between taller bars. The function returns the total amount of water trapped. There are 1 to 30,000 bars, each at most 10,000 high.
For [4, 2, 0, 3, 2, 5] the answer is 9. The bars of height 4 and 5 form walls, and between them the water fills to level 4 above the 2, level 4 above the 0, level 4 above the 3 and level 4 above the 2: that is 2 + 4 + 1 + 2 = 9 units. The standard example [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1] traps 6.
Approach: water above a bar depends on the tallest walls on each side
Think about one bar at a time. Water above bar i can rise only as high as the shorter of two walls: the tallest bar to its left and the tallest bar to its right. Any higher and it would spill over the lower wall. The water over bar i is that level minus the bar's own height.
So the answer is the sum, over all bars, of min(leftMax[i], rightMax[i]) - height[i], where leftMax[i] is the tallest bar from the first one up to bar i and rightMax[i] is the tallest from bar i to the last one. Both maxima include bar i itself, which guarantees the difference is never negative: if bar i is the tallest in its neighbourhood, the result is 0.
The direct way to get each maximum is to scan left and right from every bar, which costs O(n) per bar and O(n²) overall. The dynamic programming idea removes the repeated work. Each prefix maximum builds on the previous one: leftMax[i] is the larger of leftMax[i-1] and height[i]. Filling the slice takes one pass, and the same trick from the right end gives rightMax. A third pass adds up the water.
Check it on [4, 2, 0, 3, 2, 5]. The prefix maxima are [4, 4, 4, 4, 4, 5] and the suffix maxima are [5, 5, 5, 5, 5, 5]. Taking the minimum at each position gives [4, 4, 4, 4, 4, 5], and subtracting the heights leaves [0, 2, 4, 1, 2, 0], which adds up to 9.
Step by step
- Store the number of bars in
n. Every loop below runs over the bars. - Return 0 when there are no bars. The next lines read
height[0], which needs at least one bar. - Create
leftMaxwithmake([]int, n)and set its first element toheight[0]. Nothing comes before the first bar, so the tallest bar up to it is the bar itself. - Loop from index 1 up to
n - 1to fill the rest, since index 0 is already set. - Set each entry to the larger of the previous entry and the current bar. This reuse of the previous entry is the dynamic programming step.
- Create
rightMaxin the same way and set its last element toheight[n-1]. - Loop from
n - 2down to 0, the mirror image of the first loop. - Set each entry to the larger of the next entry and the current bar.
- Create
resultand start it at 0. - Visit every bar with a third loop.
- Add the water over the bar: the smaller of the two maxima, minus the bar's height.
- Return
result.
The solution in Go
func trap(height []int) int {
n := len(height)
if n == 0 {
return 0
}
leftMax := make([]int, n)
leftMax[0] = height[0]
for i := 1; i < n; i++ {
leftMax[i] = max(leftMax[i-1], height[i])
}
rightMax := make([]int, n)
rightMax[n-1] = height[n-1]
for i := n - 2; i >= 0; i-- {
rightMax[i] = max(rightMax[i+1], height[i])
}
result := 0
for i := 0; i < n; i++ {
result += min(leftMax[i], rightMax[i]) - height[i]
}
return result
}
Go details worth a closer look:
- Built-in
minandmax. Since Go 1.21 these are built-in functions that work on any ordered type, includingint, so no helper ormathimport is needed. make([]int, n)creates a slice of lengthnfilled with zeros. Assigning toleftMax[0]and then to later indexes is fine; unlikeappend, indexing never grows a slice, so the length must be right from the start.- The loop headers. Go has only
for. The second loop counts down withi--and testsi >= 0.iis a signedint, so it can safely reach -1 and stop the loop. - Slice parameters.
heightis a slice, a small header pointing at the caller's array. The function only reads it, so the input is left untouched. - Integer range. The result is at most 30,000 × 10,000, which is 300 million, and Go's
intis 64 bits here, so there is no overflow.
Complexity
Time: O(n). Three passes over the bars, each doing constant work per bar.
Space: O(n). The two slices leftMax and rightMax hold n integers each. A two-pointer version can get this down to O(1), but this lesson keeps the two tables because they make the idea easy to see.
Edge cases and common mistakes
- A single bar.
[5]traps 0: both maxima equal the bar, and the difference is 0. - Flat or monotonic maps.
[3, 3, 3],[1, 2, 3, 4]and[4, 3, 2, 1]all give 0, because water would run off the lower end. - A single dip.
[2, 0, 2]gives 2: the level is 2 above a bar of height 0. - Using only one maximum. Subtracting the height from
leftMax[i]alone ignores a lower wall on the right, and the sum comes out too big. - Starting the loops at the wrong index. The prefix loop starts at 1 because index 0 is already set, and the suffix loop starts at
n - 2for the same reason. Starting at 0 would readleftMax[-1], which panics. - Counting the bar itself as water. Both maxima include the bar, so subtracting
height[i]is what removes the solid part of the column.
Testing
[0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]gives 6.[4, 2, 0, 3, 2, 5]gives 9.[5]gives 0.[2, 0, 2]gives 2.[1, 2, 3, 4]gives 0.
