Count the accumulated rainfall in a 2D mountain range

Question

Challenge:

Count how much rain will be collected in the valleys between the peaks of a mountain range in a 2D world.

The mountain range is defined by a single array of positive ints, where each element defines the height a peak. IE the array [1,2,1,2] would create the 2D slice:

 . m . m
 m m m m

m = mountain
. = open space

As the rain falls, it will fill up to the height of its surrounding peaks, then overflow. If the rain reaches the edge of the grid, it disappears. In this case the amount of water collected will be 1:

 . m r m
 m m m m
 r = rain

More examples:

Input: [2, 4, 2, 2, 5, 1, 2, 2, 1, 3]

Output: 10

2D visualization:

. . . . m . . . . .
. m r r m . . . . .
. m r r m r r r r m
m m m m m r m m r m
m m m m m m m m m m

Input: [1, 2, 3, 2, 1]

Output: 0

2D visualization:

. . m . . 
. m m m .
m m m m m

Constraints:

• The input array will have size at least 1
• Each element of the array will be greater than or equal to 1
• Must be in Java 7

My Code:

//Time Complexity: worst case O(n^2)?
//Space Complexity: O(1)
public class Solution
    public static int solution(int[] heights) { 
       int fillCount = 0, minHeight = heights[0];

       for(int i = 1; i<heights.length; i++){
           if(heights[i] < minHeight){
               minHeight = heights[i];
           }
           else{
               minHeight=heights[i];
               int fillHeight = findFillHeight(heights, i); 
               if(fillHeight == -1) continue;
               fillCount+= fillIn(heights, fillHeight, i);
           }
       }

       return fillCount;
    }

    static int findFillHeight(int[] heights, int rightEdgePosition){
        int maxHeight = -1, j=0;

        for(int i=rightEdgePosition-1; i-1 >=0 && heights[i] < heights[rightEdgePosition]; i--, j=i){
            if(heights[i] >= heights[rightEdgePosition]){
                return heights[rightEdgePosition];
            }
            if(heights[i] > maxHeight){
                maxHeight = heights[i];
            }
        }

        if(heights[j] >= heights[rightEdgePosition]){
            maxHeight = heights[rightEdgePosition];
        }

        return maxHeight;
    }

    static int fillIn(int[] heights, int fillInHeight, int startHeightPosition){
        int fillCount = 0;
        for(int i=startHeightPosition-1; heights[i] < fillInHeight; i--){
            fillCount+= fillInHeight-heights[i];
            heights[i]=fillInHeight;
        }
        return fillCount;
    }
}

The idea is to search the array from left to right, looking for rises in elevation. When we find one, we backtrack looking for how high up we can fill the valley, and then redo the backtrack, this time actually filling it (this prevents "double filling" in future elevation rises).

The worst case arises when we get a large elevation on element 1 and elements 2 through n are rising from 1 to the value of element 1. IE [10, 1, 2, 3.... 10]. In this case we do the maximum number of backtracks. If I have my math right, our time complexity would be

$$n + \sum_{i=0}^{n} {i}\ = n+\frac{1}{2}n(n+1) \implies O(n^2)$$

So quadratic time obviously a red flag, but I'm not sure I can avoid it since I'm dealing with 2D arrays. Other than that, the big ugly part is the j variable in findFillHeight(). The point of it is to check that the left edge we end the loop on is properly checked. The last notable issue is that much of this code would be very risky if run on its own because the assumption is that all inputs are valid. This was due to time constraints though, so I don't consider it an issue.

Answer 1

It's possible to do it in \$O(n)\$ using \$O(n)\$ additional memory. For that we should find the amount of rain collected at each index in \$O(1)\$. It can be calculated from the heights of the highest mountain to the right and to the left:

$$\tt{rainfall[i]} = \max(0, \min(\tt{leftHighest}, \tt{rightHighest}) - \tt{height[i]})$$

We can have the height to the left updated on the fly in the same loop, but we have to precalculate the heights to the right in a separate loop.

Here is the sample code:

import java.util.Arrays;

public class Solution {
    public static int getCollectedRainfall(int[] heights) {
        int[] highestToTheRight = Arrays.copyOf(heights, heights.length + 1);
        for (int i = heights.length - 1; i >= 0; --i) {
            highestToTheRight[i] = Math.max(highestToTheRight[i], highestToTheRight[i + 1]);
        }

        int collectedRainfall = 0;
        int highestToTheLeft = 0;
        for (int i = 0; i < heights.length; ++i) {
            collectedRainfall += 
                Math.max(0, Math.min(highestToTheLeft, highestToTheRight[i]) - heights[i]);
            highestToTheLeft = Math.max(highestToTheLeft, heights[i]);
        }

        return collectedRainfall;
    }
}