42. Trapping Rain Water

Problem

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.

https://leetcode.com/problems/trapping-rain-water/

Example 1:

case1

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 * 10⁴
0 <= height[i] <= 10⁵

Test Cases

1 2	class Solution: def trap(self, height: List[int]) -> int:

solution_test.py

import pytest

from solution import Solution


@pytest.mark.parametrize('height, expected', [
    ([0,1,0,2,1,0,1,3,2,1,2,1], 6),
    ([4,2,0,3,2,5], 9),
])
@pytest.mark.parametrize('sol', [Solution()])
def test_solution(sol, height, expected):
    assert sol.trap(height) == expected

Thoughts

对于任意位置 i（0 ≤ i < n），记 l_max(i) = max{height[:i+1]}（i 及其左边最高的 bar），r_max(i) = max[height[i:]]（i 及其右边最高的 bar）。取 bound = min{l_max(i), r_max(i)}，显然 bound ≥ height[i]。如果 bound > height[i]，那么高出来的部分（bound - height[i]）都是可以有水的。

先从右向左遍历 height，记录所有的 r_max(i)。然后从左向右遍历，一边更新 l_max(i) 一边计算位置 i 的积水量，并累计总积水量。

时间复杂度 O(n)，空间复杂度 O(n)。

Code

solution.py

class Solution:
    def trap(self, height: list[int]) -> int:
        max2 = lambda a, b: a if a >= b else b
        min2 = lambda a, b: a if a <= b else b

        n = len(height)
        r_maxes = [0] * n
        r_maxes[-1] = height[-1]
        for i in range(n - 2, -1, -1):
            r_maxes[i] = max2(r_maxes[i + 1], height[i])

        total = 0
        l_max = height[0]
        for i in range(1, n - 1):
            l_max = max2(l_max, height[i])
            bound = min2(l_max, r_maxes[i])
            if bound > height[i]:
                total += bound - height[i]

        return total