Problem Statement in English

Given a 2D height map represented by a grid of integers, where each integer represents the height of the terrain at that point, the goal is to find the volume of water that can be trapped after raining.

Approach

The idea here is that if we keep processing the lowest height cell on the boundary of the unprocessed area, we can determine how much water can be trapped at that cell, since the limiting factor of its height is the minimum height of the cells surrounding it. And since we guarantee that we process the lowest height cell first, we can calculate this.

We use a min-heap (priority queue) to always process the lowest height cell first.

Solution in Python


class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        rows = len(heightMap)
        cols = len(heightMap[0])

        moves = [(-1, 0), (0, 1), (1, 0), (0, -1)]

        trapped = 0

        # all cells that are reachable from the boundary
        minheap = []

        for i in range(cols):
            # first row
            heappush(minheap, (heightMap[0][i], 0, i))
            heightMap[0][i] = -1
            # last row
            heappush(minheap, (heightMap[rows - 1][i], rows - 1, i))
            heightMap[rows - 1][i] = -1

        for i in range(rows):
            # first col
            heappush(minheap, (heightMap[i][0], i, 0))
            heightMap[i][0] = -1
            # last col
            heappush(minheap, (heightMap[i][cols - 1], i, cols - 1))
            heightMap[i][cols - 1] = -1

        maxh = -1

        while minheap:
            h, x, y = heappop(minheap)
            maxh = max(maxh, h)
            trapped += maxh - h

            for dx, dy in moves:
                nx, ny = x + dx, y + dy

                if 0 <= nx < rows and 0 <= ny < cols and heightMap[nx][ny] != -1:
                    heappush(minheap, (heightMap[nx][ny], nx, ny))
                    heightMap[nx][ny] = -1

        return trapped

Complexity

  • Time: $O(m * n * \log(m * n))$
    Here, m is the number of rows and n is the number of columns in the height map. We may need to process each cell once, and each operation on the min-heap takes $O(log(m*n))$ time.

  • Space: $O(m*n)$
    Since we are using a min-heap to store the boundary cells, in the worst case, we may end up storing all the cells in the heap.

Mistakes I Made

Had to look this one up :(

And we are done.