Learn Python implementations of the capturing rainwater interview question.
A common interview question involving arrays is the “capturing rainwater” problem (also referred to as the “trapping rainwater” problem).
Imagine a very heavy rainstorm over a highway that has many potholes and cracks. The rainwater will collect in the empty spaces in the road, creating puddles. Each puddle can only be as high as the road around it, as any excess water will just flow away.
The capturing rainwater problem asks you to calculate how much rainwater would be trapped in the empty spaces in a histogram (a chart which consists of a series of bars). Consider the following histogram:
This can be represented in Python as an array filled with the values
[4, 2, 1, 3, 0, 1, 2]. Imagine that rainwater has fallen over the histogram and collected between the bars. Here’s how the previous histogram would look filled with water:
Like with the road, the amount of water that can be captured at any given space cannot be higher than the bounds around it. To solve the problem, we need to write a function that will take in an array of integers and calculate the total water captured. Our function would return
6 for the histogram above. There are multiple ways to solve this problem, but we are going to focus on a naive implementation and an optimized implementation.
The foundation to all the solutions for this problem is that the amount of rainwater at any given index is the difference between the lower of highest bars to its left and right and the height of the index itself:
water_at_index = min(highest_left_bound, highest_right_bound) - height_of_index
Look at the histogram again. The amount of water at index
2. This is because its highest left bound is
3 (element at index
3), and its highest right bound is
2 (element at index
6). The lower of these two values is
2, and when we subtract the index’s height of
0, we get our answer of
The Naive Solution
The naive solution to the problem is to:
- Traverse every element in the array
- Find the highest left bound for that index
- Find the highest right bound for that index
- Take the lower of those two values
- Subtract the height of that index from that minimum
- Add the difference to the total amount of water
In Python, this solution looks like this:
def naive_solution(heights): total_water = 0 for i in range(1, len(heights) - 1): left_bound = 0 right_bound = 0 # We only want to look at the elements to the left of i, which are the elements at the lower indices for j in range(i+1): left_bound = max(left_bound, heights[j]) # Likewise, we only want the elements to the right of i, which are the elements at the higher indices for j in range(i, len(heights)): right_bound = max(right_bound, heights[j]) total_water += min(left_bound, right_bound) - heights[i] return total_water
While this is a functional solution, it requires nested
for loops, which means it has a big O runtime of
O(n^2). Let’s look at a solution with a more efficient runtime.
The Optimized Solution
The previous solution had a quadratic runtime, but it’s possible to solve this problem in
O(n) time by using two pointers. The pointers will start at each end of the array and move towards each other. The two-pointer approach is a common approach for problems that require working with arrays, as it allows you to go through the array in a single loop and without needing to create copy arrays.
We’ll start by creating the following variables:
total_water = 0 left_pointer = 0 right_pointer = heights.length - 1 left_bound = 0 right_bound = 0
right_pointer will start at the beginning and end of the array, respectively, and move towards each other until they meet. The algorithm is as follows:
while left_pointer < right_pointer if the height at left_pointer <= the height at right_pointer - update the left_bound to be the greater value between the current left_bound and the height at left_pointer - increment total_water to be the difference between left_bound and the height at left_pointer - move left_pointer forward by one else - update the right_bound to be the greater value between the current right_bound and the height at right_pointer - increment total_water to be the difference between right_bound and the height at right_pointer - move right_pointer back by one return total_water