LESSON 371 OF 442 · DYNAMIC PROGRAMMING

Trapping Rain Water in Scala: Dynamic Programming with Two Arrays

Given n non-negative integers representing an elevation map where the width of each bar is 1, compute how much water it can trap after raining.

Example 1:

Input: height = [0,1,0,2,1,0,1,3,2,1,2,1]
Output: 6
Explanation: The above elevation map (black section) is represented by array [0,1,0,2,1,0,1,3,2,1,2,1]. In this case, 6 units of rain water (blue section) are being trapped.

Example 2:

Input: height = [4,2,0,3,2,5]
Output: 9

Constraints:

n == height.length
1 <= n <= 3 * 10⁴
0 <= height[i] <= 10⁴

STEP 1 OF 12

Guided

Ctrl + Enter to check

SCALA · DYNAMIC PROGRAMMING

Trapping Rain Water in Scala, explained

Problem overview

The method trap receives height, an Array[Int] that describes an elevation map: each element is the height of a bar of width 1. After rain, water collects in the dips between taller bars. The method returns the total amount of water that stays, as an Int.

For [4, 2, 0, 3, 2, 5] the answer is 9. The bar of height 4 on the left and the bar of height 5 on the right form walls. Between them, the water level is limited by the lower wall, 4, so the column over the bar of height 2 holds 2 units, the column over 0 holds 4, over 3 holds 1, and over the next 2 holds 2. That adds up to 9. For [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1] the answer is 6.

There are between 1 and 30,000 bars, each between 0 and 10,000 high.

Approach: water over each bar, from the walls around it

Think about one bar at a time. Water above bar i can rise only as high as the lower of two walls: the tallest bar anywhere to its left (bar i included) and the tallest bar anywhere to its right (bar i included). Any higher and the water would spill over the lower wall. The bar itself takes up the bottom of the column, so the water above it is

min(tallestLeft, tallestRight) - height(i)

Adding this up over all bars gives the answer. Both maxima include bar i, so the value is never negative: a bar that is itself the tallest on one side holds 0.

The brute-force way computes the two maxima for every bar by scanning left and right, which is O(n²). Notice that the maxima overlap heavily. The tallest bar up to i is either the tallest up to i - 1 or bar i itself. Storing those answers in an array, so each one is built from the one before it, is dynamic programming, and it turns the scans into a single pass in each direction.

The code uses two arrays, leftMax and rightMax, then one more pass that adds the water over each bar.

Step by step

Store the number of bars in val n, since all the loops use it.
Return 0 for no bars. The next lines read height(0), which needs at least one bar. The constraints promise one, but the guard keeps an empty array from throwing. The early exit uses return.
Create leftMax with Array.fill(n)(0) and set leftMax(0) to height(0). Nothing comes before the first bar, so the tallest bar up to it is itself.
Loop over i <- 1 until n. Index 0 is already filled.
Fill each entry with math.max(leftMax(i - 1), height(i)).
Create rightMax the same way and set its last entry, rightMax(n - 1), to height(n - 1).
Loop from right to left with i <- n - 2 to 0 by -1. The by -1 makes the range count down, and to includes 0.
Fill each entry with math.max(rightMax(i + 1), height(i)), which mirrors the first loop.
Declare result with var, starting at 0. It grows once per bar.
Visit every bar with i <- 0 until n.
Add the water over bar i: math.min(leftMax(i), rightMax(i)) - height(i).
Return result as the last expression.

On [4, 2, 0, 3, 2, 5] the arrays come out like this, and the last column shows the water over each bar:

height   4 2 0 3 2 5
leftMax  4 4 4 4 4 5
rightMax 5 5 5 5 5 5
water    0 2 4 1 2 0   (sum 9)

The solution in Scala

class Solution {
    def trap(height: Array[Int]): Int = {
        val n = height.length
        if (n == 0) return 0
        val leftMax = Array.fill(n)(0)
        leftMax(0) = height(0)
        for (i <- 1 until n) {
            leftMax(i) = math.max(leftMax(i - 1), height(i))
        }
        val rightMax = Array.fill(n)(0)
        rightMax(n - 1) = height(n - 1)
        for (i <- n - 2 to 0 by -1) {
            rightMax(i) = math.max(rightMax(i + 1), height(i))
        }
        var result = 0
        for (i <- 0 until n) {
            result += math.min(leftMax(i), rightMax(i)) - height(i)
        }
        result
    }
}

The Scala details worth knowing:

Array.fill(n)(0) builds an array of n elements, each 0. The first parameter list is the length and the second is the value. The result is an Array[Int], a mutable, fixed-size array, so leftMax(i) = ... assigns into it. Because the array is held in a val, the variable can't be pointed elsewhere, but the elements can change.
Indexing and assignment use parentheses. height(i) reads and leftMax(i) = x writes. An index outside 0 until n throws ArrayIndexOutOfBoundsException.
Ranges. 1 until n stops before n. n - 2 to 0 by -1 includes both ends and counts down. If n is 1, the start n - 2 is -1, and the range -1 to 0 by -1 is empty, so the loop does nothing, as it should for one bar.
math.max and math.min work on two numbers. With two Ints they return an Int.
result += on a var. It is shorthand for result = result + .... A val would not allow it.
return and the last expression. The early return 0 leaves the method. The normal result is the final result, with no return keyword.
Int range. Int is 32 bits and wraps silently on overflow. The most water possible is below 30,000 * 10,000 = 300,000,000, which fits.

Complexity

Time: O(n). There are three passes over the bars, and each step does a constant amount of work.

Space: O(n). The arrays leftMax and rightMax each hold one number per bar. A two-pointer version can cut this to O(1) by tracking the maxima as two indices move inward, but the two arrays make the idea easier to see.

Edge cases and common mistakes

One bar, or two. A single bar has no walls and holds 0. With two bars there is nothing between them, so the answer is 0 as well.
Ascending or descending bars. [1, 2, 3] and [3, 2, 1] give 0. In both cases, on every bar, one of the two maxima is the bar itself, so the difference is 0.
A flat valley. [3, 0, 3] gives 3: both walls are 3, and the middle bar is 0.
Using the global maximum for both sides. The water level over a bar depends on the walls on its own sides, not on the tallest bar in the whole map.
Leaving the bar out of the maxima. If a maximum didn't include bar i, the difference could be negative for a tall bar, and the sum would be wrong. Starting leftMax(0) from height(0) and taking math.max with height(i) keeps the bar in.
Off-by-one in the right-to-left loop. The loop starts at n - 2 because the last entry is already set. Starting at n - 1 would read rightMax(n), one past the end of the array, and throw.
Forgetting by -1. n - 2 to 0 with the default step is empty when n - 2 > 0, so rightMax would stay zero and the result would be wrong.

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.
[3, 0, 3] gives 3.
[1, 2, 3] gives 0.
[3, 2, 1] gives 0.

Keep going with ScalaAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free