Hard

Trapping Rain Water II

LAB

Given a height map, return how much rain water can be trapped in the 2D terrain.

EXAMPLES

Example 1

Input
{
  "heightMap": [
    [
      1,
      4,
      3,
      1,
      3,
      2
    ],
    [
      3,
      2,
      1,
      3,
      2,
      4
    ],
    [
      2,
      3,
      3,
      2,
      3,
      1
    ]
  ]
}

Output
4

FUNCTION SHAPE

heightMap: intMatrix→int

SOLUTION NOTE

This is actually a really clever usage of dijkstras. Inspired by trapping rain water I, where we needed to min max_prefix, max_suffix paths, here we need the min max values for all paths from the border to that cell.

Basically, for each single cell, we need to know that for all the possible paths to the outside world (where the water will escape to), what is the minimum of all path's weight, and the path's weight should be defined as the highest height value along the path.

The naive approach is to just consider the 4 directions, up, left, down, right in terms of max prefix paths. However, this doesn't work because water paths may zigzag. Imagine water flowing on the grid from the sky, water will flow from the borders and collect in the valleys. For a particular valley/cell, we want to consider all possible paths water could have flowed from the border, and for each path, consider their max height that they encountered on that path and take the minimum max height over all possible paths. This determines the height of the trapped water for this cell, and we simply need to subtract from the height of this cell to get the trapped water amount.

How can we implement finding the minimum max height path out of all paths to reach a certain node? This is precisely what dijkstra's gives us. Dijkstra's guarantees the first time we reach a node, it is with the minimum path.

Guidance: It's a good idea for all problems where you think dijkstra's is applicable, to quickly think about that cloud proof and that counterexample case for 30s-1minute. This will show the interviewer you are rigorous with correctness and not just purely intuitive. Also, this will give you the confidence to invest time into implementing the approach.

Reveal reference solution +

pythonREFERENCE

def trapRainWater(self, A: List[List[int]]) -> int:
    m,n = len(A), len(A[0])
    heap = []
    border = set()
    for i in range(m):
        for j in range(n):
            if i in [0, m-1] or j in [0, n-1]:
                border.add((i,j))
                heappush(heap, (A[i][j], i, j))

    dirs = [(0,1),(1,0),(-1,0),(0,-1)]
    def isInBounds(i,j):
        return 0 <= i < m and 0 <= j < n

    res = 0
    d = [[inf for _ in range(n)] for _ in range(m)]

    while heap:
        max_height_so_far, i, j = heappop(heap)

        if d[i][j] < max_height_so_far: continue

        for x,y in dirs:
            ii,jj = i+x,j+y
            if isInBounds(ii,jj) and (ii,jj) not in border and (dd := max(max_height_so_far, A[ii][jj])) < d[ii][jj]:
                d[ii][jj] = dd

                res += max(0, max_height_so_far - A[ii][jj])
                heappush(heap, (dd, ii, jj))

    return res

TimeO(m*n*log(m*n))

SpaceO(m*n)

Open on LeetCode ↗