Rain Water Trapping Problem

Difficulty: Hard, Asked-in: Google, Microsoft, Amazon

Key takeaway: This is an excellent problem to learn time and space complexity optimization using various approaches. Two-pointers and stack-based approaches are worth exploring.

Let’s understand the problem

Given n non-negative integers representing an elevation map where the width of each tower is 1. Write a program to compute how much water it can trap after raining.

Example

Input: height[] = [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1], Output: 6

Explanation: Trapped water = 1 x 1 + 1 x 1 + 2 x 1 + 1 x 1 + 1 x 1 = 6 (Area of the blue region in the following diagram)

Discussed solution approaches

• Brute force solution using nested loops
• Using single loop and extra memory
• Using stack data structure
• Efficient approach using two pointers

Brute force solution  using nested loops

Solution idea

One basic idea would be to traverse the array, calculate the rainwater trapped at each tower and do the sum to get the overall water trapped. The critical question is: How do we find the amount of water trapped at each tower height[i]?

If we observe closely, water trapped by tower height[i] will be bounded by the maximum height on the left and right sides of height[i]. So we traverse elements on the left of height[i] to find the highest tower on the left side. Similarly, we traverse elements on the right of height[i] to find the highest tower on the right side. Note: We need to include the current height in the calculation for finding the max height on the left and right sides. Why? Think!

Finally, total amount of water trapped at height[i] will be the difference between smaller heights and the height of the current element height[i]. Rainwater trapped at the height[i] = min (maximum height on the left side, maximum height on the right side) - height[i].

Solution steps

1. We initialize variable trappedWater to store the total trapped water.
2. Now we traverse array from i = 0 to n - 1. Inside the loop, we initialize variables leftMax and rightMax to track the maximum height of towers on both sides. For each height[i]:

3. We run a loop from j = i to 0 to find maximum height on the left side and track this value in variable leftMax.

   int leftMax = 0
   for(int j = i; j >= 0; j = j - 1)
     leftMax = max(height[j], leftMax)

4. Similarly, we run a loop from j = i to n - 1 to find maximum height on the right side and track this value in variable rightMax.

   int rightMax = 0
   for(int j = i; j < n; j = j + 1)
     rightMax = max(height[j], rightMax)

5. Now using the formula, we calculate the amount of water trapped at height[i] and add it to the variable trappedWater.

   trappedWater = trappedWater + min(leftMax, rightMax) - height[i]

6. By the end of nested loops, we return value stored in variable trappedWater.

Solution code C++

int rainWaterTrapping(int height[], int n)
{
    if (n <= 2)
        return 0;
    int trappedWater = 0;
    for (int i = 0; i < n; i = i + 1)
    {
        int leftMax = 0;
        for (int j = i; j >= 0; j = j - 1)
            leftMax = max(height[j], leftMax);

        int rightMax = 0;
        for (int j = i; j < n; j = j + 1)
            rightMax = max(height[j], rightMax);

        trappedWater = trappedWater + min(leftMax, rightMax) - height[i];
    }
    return trappedWater;
}

Solution code Python

def rainWaterTrapping(height, n):
    if n <= 2:
        return 0
    trapped_water = 0
    for i in range(n):
        left_max = 0
        for j in range(i, -1, -1):
            left_max = max(height[j], left_max)

        right_max = 0
        for j in range(i, n):
            right_max = max(height[j], right_max)

        trapped_water = trapped_water + min(left_max, right_max) - height[i]
    return trapped_water

Solution analysis

We use two nested loops where outer loop is scanning height[] array, and two inner loops are finding rightMax and leftMax. So at each iteration of outer loop, we are traversing each array element using inner loops. So, time complexity = n * O(n) = O(n^2).

We are using constant number of variables. So, space complexity = O(1)

Using single loop and extra memory

Solution idea

Critical questions are: Can we improve efficiency and solve the problem in O(n) time? Is it possible to get rid of two inner loops? Let's think!

If values of leftMax and rightMax are known for each tower height[i], we can solve this problem in a single scan of array. How can we do this? One idea would be to do some pre-processing, calculate both values for each tower and store them in two separate extra memory of size n.

Suppose we take two extra arrays, leftMax[n] and rightMax[n] i.e. leftMax[i] to store max height on the left side and rightMax[i] to store max height on the right side. Can we calculate both arrays in O(n) time? Here is an idea:

• If we know leftMax[i - 1], we can easily calculate leftMax[i] in O(1) using the formula leftMax[i] = max (leftMax[i - 1], height[i]).
• Similarly, if we know rightMax[i + 1], we can easily calculate rightMax[i] in O(1) using the formula rightMax[i] = max (rightMax[i + 1], height[i]).

We can store leftMax[] and rightMax[] arrays in a single scan of the height[] array. This is a dynamic programming approach where we use the already calculated solution of the smaller problem to get the solution to the larger problem. Think!

Solution steps

• We create two arrays leftMax[n] and rightMax[n].

• Now we initialize leftMax[0] with height[0] and run a loop from i = 1 to n - 1 to store values in leftMax[]. At each ith iteration, we store max element that occurred up to that point in leftMax[i].

  leftMax[0] = height[0]
  for(int i = 1; i < n; i = i + 1)
    leftMax[i] = max(leftMax[i - 1], height[i])

• Similarly, we initialize rightMax[n - 1] with height[n - 1] and run a loop from i = n - 2 to 0 to store values in rightMax[]. At each ith iteration, we store max element that occurred up to that point in rightMax[i].

  rightMax[n - 1] = height[n - 1]
  for(int i = n - 2; i >= 0; i = i - 1)
    rightMax[i] = max(rightMax[i + 1], height[i])

• Stored water at height[i] = min(leftMax[i], rightMax[i]) - height[i]. So we traverse height[] array and track total amount of water trapped using the variable trappedWater.

  int trappeWater = 0
  for(int i = 0; i < n; i = i + 1)
    trappedWater = trappedWater + min(leftMax[i], rightMax[i]) - height[i]

• Finally, we return value stored in the variable trappedWater.

Solution code C++

int rainWaterTrapping(int height[], int n)
{
    if (n <= 2)
        return 0;
    int leftMax[n], rightMax[n];

    leftMax[0] = height[0];
    for (int i = 1; i < n; i = i + 1)
        leftMax[i] = max(leftMax[i - 1], height[i]);

    rightMax[n - 1] = height[n - 1];
    for (int i = n - 2; i >= 0; i = i - 1)
        rightMax[i] = max(rightMax[i + 1], height[i]);

    int trappedWater = 0;
    for (int i = 0; i < n; i = i + 1)
        trappedWater = trappedWater + min(leftMax[i], rightMax[i]) - height[i];

    return trappedWater;
}

Solution code Python

def rainWaterTrapping(height, n):
    if n <= 2:
        return 0
    left_max = [0] * n
    right_max = [0] * n

    left_max[0] = height[0]
    for i in range(1, n):
        left_max[i] = max(left_max[i - 1], height[i])

    right_max[n - 1] = height[n - 1]
    for i in range(n - 2, -1, -1):
        right_max[i] = max(right_max[i + 1], height[i])

    trapped_water = 0
    for i in range(n):
        trapped_water = trapped_water + min(left_max[i], right_max[i]) - height[i]

    return trapped_water

Solution analysis

We are running three single loops. So, time complexity = Time complexity of storing values in leftMax[] + Time complexity of storing values in rightMax[] + Time complexity of calculating total trapped water = O(n) + O(n) + O(n) = O(n).

We are using two extra arrays of size n. So space complexity = O(n) + O(n) = O(n)

Using stack data structure

Solution idea

The critical question is: Can we solve this problem in a single scan of height[] array? Let's explore.

• Region 1 Area = Area bounded between index 2 and 0
• Region 2 Area = Area bounded between index 5 and 3
• Region 3 Area = Area bounded between index 6 and 3
• Region 4 Area = Area bounded between index 6 and 2
• Region 5 Area = Area bounded between index 8 and 6
• Region 6 Area = Area bounded between index 9 and 6

Did we observe some patterns? The idea would be to track the area confined between current tower and all previous smaller towers in a sequence. We can use a stack to track indices of the previous smaller towers. This can be done using one iteration.

• We will keep pushing elements in stack until we found an element larger than value at stack top. This means that there is a chance of storing water on the left side of current tower.
• So we remove smaller element on the left side by popping element from stack and find the amount of water stored by popped tower and current tower. We continue this till stack is not empty or a higher-value element is found.

Solution steps

• We create a stack S to store indices of towers.
• Now we traverse height[] array using a loop: while(i < n).

• If we found current tower height[i] larger than tower at the top of stack, we can say that tower at top is confined between current tower and a previous tower in the stack. So we pop value from stack and add water trapped between towers to the total water trapped. We continue this step in a loop until we find current tower smaller than tower at the top of stack or stack becomes empty.

  • Water trapped between towers = Area of the rectangular region formed by current tower, popped tower, and tower at the top of the stack.
  • Region length = Current index - Index of top stack element - 1
  • Region height = min (Height of current tower, Height of tower at the top of stack) - Height of popped tower
  • Water trapped = Region length x Region height
• If current tower is smaller than or equal to the tower at top, we push index of the current tower to stack and move to the next tower.
• Finally, we return value stored in the variable trappedWater.

Solution code C++

int rainWaterTrapping(int height[], int n)
{
    int trappedWater = 0;
    int i = 0;
    stack<int> S;
    while (i < n)
    {
        while (!S.empty() && height[i] > height[S.top()])
        {
            int topTower = S.top();
            S.pop();
            if (S.empty())
                break;
            int regionLength = i - S.top() - 1;
            int regionHeight = min(height[i], height[S.top()]) - height[topTower];
            trappedWater = trappedWater + regionLength * regionHeight;
        }
        S.push(i);
        i = i + 1;
    }
    return trappedWater;
}

Solution code Python

def rainWaterTrapping(height, n):
    trappedWater = 0
    i = 0
    S = []
    while i < n:
        while S and height[i] > height[S[-1]]:
            topTower = S.pop()
            if not S:
                break
            regionLength = i - S[-1] - 1
            regionHeight = min(height[i], height[S[-1]]) - height[topTower]
            trappedWater = trappedWater + regionLength * regionHeight
        S.append(i)
        i = i + 1
    return trappedWater

Solution code Java

class RainWater
{
    public int rainWaterTrapping(int[] height, int n) 
    {
        int trappedWater = 0;
        int i = 0;
        Stack<Integer> S = new Stack<>();
        while (i < n) 
        {
            while (!S.empty() && height[i] > height[S.peek()]) 
            {
                int topTower = S.pop();
                if (S.empty())
                    break;
                int regionLength = i - S.peek() - 1;
                int regionHeight = Math.min(height[i], height[S.peek()]) - height[topTower];
                trappedWater = trappedWater + regionLength * regionHeight;
            }
            S.push(i);
            i = i + 1;
        }
        return trappedWater;
    }
}

Solution analysis

We are doing a single traversal of the array. In the worst case, we will be performing two stack operations for each tower i.e. one push() and one pop() operation. (Think!) Both push() and pop() operations take O(1) time in the worst case. So, time complexity = O(n). In the worst case, the stack can take up to n elements, so pace complexity = O(n). 

Efficient solution using two pointers approach

Solution idea

Based on the formula used in the 2nd approach, the water trapped by any tower depends on the maximum height of towers on both sides. So instead of maintaining two arrays of size n for storing leftMax and rightMax, we can think of maintaining two variables to store the maximum heights till a given point. In other words, instead of computing the left and right parts separately, we may think of some way to do it in one iteration.

Suppose we start from both the left and right end by maintaining two pointers left and right. If the smaller tower is at the left end, water trapped would be dependent on the tower's height in the direction from left to right. Similarly, if the tower at the right end is smaller, water trapped would be dependent on the tower's height in the direction from right to left. So we first calculate water trapped on the smaller tower among height[left] and height[right] and move the pointer associated with the smaller tower.

Loop will run till the left pointer doesn't cross the right pointer. In terms of an analogy, we can think of height[left] and height[right] as a wall of a partial container where we fix the higher end and flow water from the lower end.

Solution steps

• We Initialise left and right pointers.

• We also Initialise variables trappedWater, leftMax, and rightMax.

  int trappedWater = 0
  int leftMax = 0, rightMax = 0
  int left = 0, right = n - 1

• Now we run a loop to scan the array i.e. while (left <= right)

  1. If height[left] < height[right] and height[left] < leftMax: We calculate the amount of rainwater trapped between both towers and increase the left pointer. But if height[left] > leftMax: We update the value of leftMax and increase the left pointer.

  if (height[left] < height[right])
  {
      if (height[left] > leftMax)
          leftMax = height[left]
      else
          trappedWater = trappedWater + leftMax - height[left]
      left = left + 1
  }

  2. If height[left] > height[right]) and height[right] < rightMax: We calculate the amount of rainwater trapped between both towers and decrease the right pointer. But if height[right] > rightMax: We update the value of rightMax and decrease the right pointer.

  if (height[left] > height[right])
  {
      if (height[right] > rightMax)
          rightMax = height[right]
      else
          trappedWater = trappedWater + rightMax - height[right]
      right = right - 1
  }

• When left > right, then we exit the loop and return the value stored in the trappedWater.

Solution code C++

int rainWaterTrapping(int height[], int n) 
{
    int trappedWater = 0;
    int leftMax = 0, rightMax = 0;
    int left = 0, right = n - 1;
    while (left <= right) 
    {
        if (height[left] < height[right]) 
        {
            if (height[left] > leftMax)
                leftMax = height[left];
            else
                trappedWater = trappedWater + leftMax - height[left];
            left = left + 1;
        } 
        else 
        {
            if (height[right] > rightMax)
                rightMax = height[right];
            else
                trappedWater = trappedWater + rightMax - height[right];
            right = right - 1;
        }
    }
    return trappedWater;
}

Solution code Python

def rainWaterTrapping(height, n):
    trappedWater = 0
    leftMax = 0
    rightMax = 0
    left = 0
    right = n - 1
    while left <= right:
        if height[left] < height[right]:
            if height[left] > leftMax:
                leftMax = height[left]
            else:
                trappedWater = trappedWater + leftMax - height[left]
            left = left + 1
        else:
            if height[right] > rightMax:
                rightMax = height[right]
            else:
                trappedWater = trappedWater + rightMax - height[right]
            right = right - 1
    return trappedWater

Solution analysis

We are doing a single scan to traverse the array from both ends. After each comparison, we move either the left or right pointer. So in the worst case, we need to do O(n) operations. So, time complexity = O(n).

Only constant space is required for variables and pointers, so space complexity = O(1).

Critical ideas to think!

• Can we find another solution to this problem in O(n) time and O(1) space? Here is a hint: we first search for the maximal tower in height [] and split the array into two halves around it. Now we do two traversals: one from the leftmost tower to the highest tower and another from the rightmost to the highest tower.
• In the dynamic programming approach, can we save one pass? Here is a hint: Use one array to keep track of the maximum height on one side, i.e., left or right. On the other side, we can use a variable to keep track of the maximum height so far and process it on the fly during the process when we calculate the volume of water in each tower.
• In the stack-based approach, what would be the best and worst-case scenario?
• In the 4th approach, can we compare leftMax and rightMax to decide which pointer to move?
• Let's try a different version of the problem: given an m x n integer matrix height[][] representing the height of each unit cell in a 2D elevation map, return the volume of water it can trap after rain.

Comparisons of time and space complexities

• Brute force approach: Time = O(n²), Space = O(1)
• Using dynamic programming: Time = O(n), Space = O(n)
• Using stack: Time = O(n), Space = O(n)
• Using two pointers: Time = O(n), Space = O(1)

Similar coding questions to practice

• Container with most water
• Check for pair in an array with a given sum
• Sort an array of 0s, 1s, and 2s
• Remove duplicates from sorted array
• Equilibrium index of an array

Please write comments if you find an error or you want to share more insights about the topic. Enjoy learning, Enjoy coding, Enjoy algorithms!