Leetcode 42: Trapping rain water

leetcode_42.md

Description

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.

Link to the original problem

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 <= 2 * 104
• 0 <= height[i] <= 105

Idea

1. This is another famous stack problem that I cannot miss.
2. Also it's asked frequently in interviews, so definitely worth looking at.
3. Also after we reviewed the Leetcode 85: Maximal Rectangle, this one should somehow feel similar. Let's see how.
4. So we are asked to calculate how much rain can be trapped? A silly question is, how can we trap water? The answer is we need something higher on the sides and lower in the middle. Like a bucket, right?
5. In this question's example 1, we see three small "bucket"s colored in blue, so how do we compute it?
6. Since the "buckets" are not always a rectangle or square, obviously, we have to "divide and conquer".
7. Especially, for the second "bucket", depending on how you "divide" the bucket, you'll get different solutions. Here we try to divide it horizontally so we get horizontal rectangles. Like below.
8. Why we are dividing by horizontally? It's because we want to take advantage of Stack's magic.
9. So for each horizontal rectangle in the above figure, if we can find its height and width, the problem is solved.
10. I admit, the connection with Stack is a bit tricky to find.
11. Let's try thinking of it this way. So when we iterate through each bar, we need to meet a "bucket" to be able to trap some water. Imagining we are walking on the bars in the elevation maps, literally "we need to walk downwards (like into a valley) and then upwards to be able to trap water". And also when we walk downwards, we don't know how much water we can trap, so basically, we "have" to keep going, until we walk upwards. And each step upwards will trap "some" water. This "some" water literally translates into the rectangles we divided, like in the above figure.
12. So the question is what kind of data structure will support this behavior? Stack!
13. Also if you remember what we did in Leetcode 85: Maximal Rectangle, it'll somehow reminds you to be able to think "retrospectively".
14. So the idea should be clear now. We keep recording the indices of the bars we walked by in the stack, and we try to compute the water trapped when we walk upwards (i.e., when we meet higher bars). And since we recorded the indices in the stack, it'll support us to find the height (from the index at the stack top), and width (from current index i and the index before the stack top).
15. The code should make more sense now, I guess. Let me know if otherwise.

Solution

Javascript

function trap(height) {
  let ans = 0, i = 0;
  const stk = [];
  while (i < height.length) {
    while (stk.length && height[i] > height[stk[stk.length-1]]) {
      const top = stk.pop();
      if (!stk.length) break;
      const leftI = stk[stk.length-1];
      const width = i - leftI - 1;
      const boundedHeight = Math.min(height[i], height[leftI]) - height[top];
      ans += width * boundedHeight;
    }
    stk.push(i);
    i += 1;
  }
  return ans;
}

References

1. Leetcode 85: Maximal Rectangle