Problem overview
The method trap receives height, an array of non-negative integers. Each value is the height of a bar of width 1, standing side by side. After rain, water collects in the dips between taller bars. The method returns how many units of water stay trapped.
For [4, 2, 0, 3, 2, 5] the answer is 9. Above the bar of height 2 at index 1, water rises to level 4 (the wall on the left is 4, the wall on the right is 5), which holds 2 units. Above the bar of height 0 it also rises to 4, giving 4 units. Above 3 it holds 1 unit, and above the second 2 it holds 2 units. Together: 2 + 4 + 1 + 2 = 9. The bars at both ends hold nothing, since nothing stands beyond them.
Approach: water above each bar
Think about one bar at a time. The water above bar i rises until it would spill over the lower of the two walls around it: the tallest bar to its left and the tallest bar to its right (both counting bar i itself). So the water above it is
min(tallest on the left, tallest on the right) - height[i]Adding this over all bars gives the answer. The straightforward way to find the two maxima scans left and right from every bar, which is O(n²). But the tallest bar up to i follows from the tallest up to i - 1: it is the larger of that value and height[i]. Computing it for every position in one pass from the left, and likewise from the right, takes O(n) time and gives each lookup in constant time. Storing earlier results to reuse them is the dynamic programming here. A two-pointer version needs no arrays, but this lesson uses the two arrays because they make the idea easy to see.
For [4, 2, 0, 3, 2, 5] the left maxima are [4, 4, 4, 4, 4, 5] and the right maxima are [5, 5, 5, 5, 5, 5]. The smaller of the two at each index is [4, 4, 4, 4, 4, 5], and subtracting the heights gives [0, 2, 4, 1, 2, 0], which adds up to 9.
Step by step
- Store the number of bars in
n, because every loop below uses it. - Return 0 when there are no bars. The next lines read
height[0], which needs at least one bar. Swift requires braces around the body, even for one statement. - Create
leftMaxasnzeros and set its first element toheight[0]: the tallest bar up to the first bar is the bar itself. - Loop over
1..<n. Index 0 is already filled. - Fill each entry with the larger of the previous entry and
height[i]. - Create
rightMaxthe same way and set its last element toheight[n - 1]. - Loop backward with
stride(from: n - 2, through: 0, by: -1), since the last entry is already set. - Fill each entry with the larger of the entry after it and
height[i]. - Start
resultat 0 withvar. - Loop over every bar,
0..<n. - Add the water above bar
i: the smaller ofleftMax[i]andrightMax[i], minusheight[i]. Both maxima include the bar itself, so the difference is never negative. - Return
result.
The solution in Swift
class Solution {
func trap(_ height: [Int]) -> Int {
let n = height.count
if n == 0 { return 0 }
var leftMax = [Int](repeating: 0, count: n)
leftMax[0] = height[0]
for i in 1..<n {
leftMax[i] = max(leftMax[i - 1], height[i])
}
var rightMax = [Int](repeating: 0, count: n)
rightMax[n - 1] = height[n - 1]
for i in stride(from: n - 2, through: 0, by: -1) {
rightMax[i] = max(rightMax[i + 1], height[i])
}
var result = 0
for i in 0..<n {
result += min(leftMax[i], rightMax[i]) - height[i]
}
return result
}
}
Swift details worth knowing:
- Creating an array of zeros.
[Int](repeating: 0, count: n)builds an array withnelements, which the loops then overwrite. WritingleftMax[i]before the array has that many elements would crash, because Swift checks array bounds at run time. - Ranges.
1..<nexcludesn, so it covers indices 1 ton - 1. Forn = 1it is empty and the loop does nothing. A range whose start is above its end, such as1..<0, is not empty in Swift: creating it crashes the program. That is why the check forn == 0comes first. - Counting down.
0..<ncan't run backward.stride(from:through:by:)includes the end value (through), andby: -1counts down. Forn = 1the loop starts at-1, so it makes no iterations, which is correct sincerightMax[0]is already set. minandmax. These are global functions taking two values, as inmax(leftMax[i - 1], height[i]).- Integer overflow. Swift's
Intis 64 bits and overflow stops the program rather than wrapping. Here the largest possible total is about 3 * 10^4 bars times 10^4 height, far below the limit.
Complexity
Time: O(n). There are three separate loops over the bars, each with constant work per step.
Space: O(n). The two arrays each hold n integers. The two-pointer version brings this down to O(1), at the price of a trickier invariant.
Edge cases and common mistakes
- One or two bars.
[5]traps nothing: both maxima equal the bar, so the difference is 0. Two bars can't hold water either, since there is nothing between them. - Monotonic heights.
[1, 2, 3]and[3, 2, 1]both give 0: on a staircase the water always spills over the lower side. - A single dip.
[2, 0, 2]gives 2. - Using only one side. Taking the left maximum alone, or the global maximum, overestimates:
[3, 0, 0, 2, 0, 4]has the answer 10, and a global maximum of 4 would claim a level of 4 above the 2. - Forgetting the bar itself. Because both maxima include bar
i,height[i]can never exceed their minimum. Computing the maxima over the other bars only would allow negative terms, and the sum would be wrong. - Counting water on the end bars. No wall stands beyond the first or last bar, and the formula returns 0 for them automatically, since the maximum on the outer side is the bar itself.
Testing
[0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]gives 6.[4, 2, 0, 3, 2, 5]gives 9.[5],[1, 2, 3]and[3, 2, 1]give 0.[2, 0, 2]gives 2, and[3, 0, 0, 2, 0, 4]gives 10.
