We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

Trapping Rain Water II

Number: 407

Difficulty: Hard

Paid? No

Companies: Google, Walmart Labs, Amazon, Bloomberg, ByteDance, Meta, Adobe, Oracle, X


Problem Description

Given an m x n matrix representing the height of each unit cell in a 2D elevation map, compute how much water can be trapped after it rains. Water is trapped in dips bounded by higher cells and cannot leak off the boundary.


Key Insights

  • The problem is an extension of the trapping rain water problem from 1D to 2D.
  • Use a min-heap (priority queue) to always process the lowest boundary cell.
  • Start from the boundary and gradually work inward, updating water level based on the current minimum height.
  • A cell can trap water if its height is less than the current water level; update the water level as max(current height, neighbor height).
  • Use BFS to traverse neighbor cells to ensure all cells are considered only once.

Space and Time Complexity

Time Complexity: O(m * n * log(m * n)) because each cell is pushed and popped from the heap. Space Complexity: O(m * n) due to the heap and visited flags.


Solution

We solve the problem using a multi-source breadth-first search (BFS) that originates from all boundary cells. A min-heap keeps track of the current boundary, ensuring that we always examine the cell with the smallest height. For each extracted boundary cell, we check its neighbors: if a neighbor has not been visited, it can potentially trap water. The trapped water is computed as the difference between the current boundary height and the neighbor’s height (if positive). The neighbor is then updated to have an effective height equal to the maximum of its original height and the current boundary, and it is pushed to the heap. This approach prevents overlooking lower barriers inside the grid while ensuring we only compute trapped water for valid depressions.


Code Solutions

# Import heapq for the priority queue implementation.
import heapq

def trapRainWater(heightMap):
    if not heightMap or not heightMap[0]:
        return 0
    
    m, n = len(heightMap), len(heightMap[0])
    visited = [[False] * n for _ in range(m)]
    heap = []
    
    # Add all boundary cells into the heap and mark as visited.
    for i in range(m):
        for j in range(n):
            if i == 0 or i == m - 1 or j == 0 or j == n - 1:
                heapq.heappush(heap, (heightMap[i][j], i, j))
                visited[i][j] = True

    water_trapped = 0
    directions = [(1, 0), (-1, 0), (0, 1), (0, -1)]
    
    # Process cells from the boundary inward.
    while heap:
        height, x, y = heapq.heappop(heap)
        # Explore all 4-directional neighbors.
        for dx, dy in directions:
            nx, ny = x + dx, y + dy
            if 0 <= nx < m and 0 <= ny < n and not visited[nx][ny]:
                visited[nx][ny] = True
                # If the neighbor's height is lower, water can be trapped.
                water_trapped += max(0, height - heightMap[nx][ny])
                # Update the neighbor's height to the boundary level.
                heapq.heappush(heap, (max(heightMap[nx][ny], height), nx, ny))
    return water_trapped

# Example usage:
# heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
# print(trapRainWater(heightMap))  # Output: 4
← Back to All Questions