Earlier, we tackled the Trapping Rain Water problem, which required us to determine the volume of water that an elevation map, represented by an array of non-negative integers, would be able to hold.
In this article, we’ll be looking at the 2D version of the problem.
From LeetCode:
Given an
nxmmatrix of positive integers representing the height of each cell in a 2D elevation map, compute the amount of water that can be trapped within it.Note: Both
nandmare less than 110. The height of each cell is greater than 0 and is less than 20,000.Input: [ [1, 4, 3, 1, 3, 2], [3, 2, 1, 3, 2, 4], [2, 3, 3, 2, 3, 1] ] Output: 4
I recommend working through the first version of this problem if you haven’t, because it’ll provide you with some idea for how you can solve this one.
For this question in particular, getting the intuition down will be far more helpful than knowing what snippets of code to write. I’ll save the solution (annotated with code comments) for the end.
From a 1D Array to a 2D Matrix
In Trapping Rain Water, we established that the volume of water held above a single index would be determined by the shorter of the two max-height bars located on its left and right.
We initialised two pointers at each end and moved them inward using a greedy algorithm. At each index, we identified the side of the elevation map that imposed an upper bound on the volume of water supported.
In this problem, we’re confronted with the same challenge in matrix form. One of our jobs is to find the upper bound on the amount of water contained in each cell; it doesn’t help that the water now flows in four rather than two directions! As such, we’ll need to maintain some kind of information regarding the entire perimeter of the matrix.
We can’t possibly go from using two pointer variables in 1D to using 2n + 2m pointers in 2D. However, we can store references to all of the perimeter coordinates. With these references stored somewhere, we can greedily examine each of them, one at a time.
Of course, this raises the question: Greedily? In what way? Let’s refer to an example.
Peeking from the Perimeter
[5, 5, 5, 5]
[5, ?, ?, 7]
[5, ?, ?, 3]
[7, 7, 7, 7]We’re given the above matrix. Suppose we’re only aware of the heights at the perimeter. All the perimeter cells are of height 5 or 7, with the exception of mtx[2][3] that’s at height 3.
It’s fair to say that any water body that’s trapped (i.e. blocked) by mtx[2][3] will also be trapped by the other perimeter cells, since they are of greater height. For example:
[5, 5, 5, 5]
[5, ?, ?, 7]
[5, ?, 2, 3]
[7, 7, 7, 7]
# mtx[2][2] supports 1 unit of water and there's no
# way this unit of water can escape, since mtx[2][3]
# is lower than all the other perimeter cells.The inverse doesn’t hold true:
[5, 5, 5, 5],
[5, 4, ?, 7],
[5, ?, ?, 3],
[7, 7, 7, 7],
# mtx[1][1] *might* hold 1 unit of water, but this
# unit of water *might* also flow out via mtx[2][3].In a sense, we’re more confident about what we can say about the inner cells adjacent to mtx[2][3] than what we can say about the inner cells adjacent to other perimeter cells. Thus, we’ll traverse the matrix with mtx[2][3] as our starting point.
Given that mtx[2][3] is of height 3, it can trap a water body composed of cells whose heights are all lower than 3. By extension, each of the cells in these water bodies can hold at most 3 units of water.
We can also say that 3 is our current upper bound on the amount of trappable water. It’s important that we keep track of this value.
Let’s consider what happens when we traverse the matrix from mtx[2][3].
Navigating the Matrix
We move left from mtx[2][3], half-expecting to find a cell of lesser height, but we find that mtx[2][2] is of height 9!
# Before
[5, 5, 5, 5]
[5, ?, ?, 7]
[5, ?, ?, 3]
[7, 7, 7, 7]
# After
[5, 5, 5, 5]
[5, ?, ?, 7]
[5, ?, 9, 3]
[7, 7, 7, 7]Clearly, no water can be trapped at mtx[2][2] since it’s taller than our current upper bound of 3.
What’s next? Taking into account at the rest of the grid that we’ve visited, we’ll want to traverse the matrix again from the next lowest cell we’ve seen thus far. Specifically, we could continue our traversal from any one of the perimeter cells with a height of 5 –– for instance, mtx[0][1].
Why? To retread the point made earlier, we know that any water body that’s trapped by mtx[0][1] will be trapped by the other cells we’ve seen (including mtx[2][2]) since they are of greater or equal height.
Our current upper bound on the amount of trappable water was previously 3 based on mtx[2][3]. Now, we’ll update this value to 5 because we’ve made mtx[0][1] our new point of entry into the matrix.
We walk from mtx[0][1] to mtx[1][1] to find a cell with a height of 2.
# Before
[5, 5, 5, 5]
[5, ?, ?, 7]
[5, ?, 9, 3]
[7, 7, 7, 7]
# After
[5, 5, 5, 5]
[5, 2, ?, 7]
[5, ?, 9, 3]
[7, 7, 7, 7]How much water can be contained at mtx[1][1]? We subtract 2 from our current upper bound, 5, which gives us 3 units of water.
At this point, of all the cells we’ve seen, mtx[1][1] lies the lowest. We’ll continue exploring from this position since the water body bounded by mtx[0][1] and all the other cells might extend further.
# Before
[5, 5, 5, 5]
[5, 2, ?, 7]
[5, ?, 9, 3]
[7, 7, 7, 7]
# After
[5, 5, 5, 5]
[5, 2, 4, 7]
[5, 0, 9, 3]
[7, 7, 7, 7]We find mtx[1][2] with a height of 4 (which yields us 1 unit of water) and mtx[2][1] with a height of 0 (which yields us 5 units of water).
In total, this grid is able to trap 9 units of water.
Data Structure Requirements
Two requirements become quite apparent to us:
- We’ll need to greedily pick the lowest cells we’ve seen as points for traversing the matrix
- We’ll need keep track of the upper bound on the amount of trappable water when exploring each trapped water body
To achieve the first, we’ll want to store our seen cells in a min heap:
Cell = collections.namedtuple('Cell', ('ht', 'r', 'c'))
# ...
pq = []
# ...
heapq.heappush(pq, Cell(mtx[r][c], r, c))Tracking the upper bound is quite straightforward –– we simply can use an ordinary variable:
popped_max = 0
# ...
popped = heapq.heappop(pq)
popped_max = max(popped_max, popped.ht)Review
To see the full algorithm in action, I highly recommend checking the following visualisation by Sam Shen:
Given nm cells, our algorithm runs in O(nm) space and O(nm log nm) time.
With that, we’ve solved the Trapping Rain Water (2D) problem 🌧📥.
Full Solution
import heapq
import collections
Cell = collections.namedtuple('Cell', ('ht', 'r', 'c'))
SHIFT = [(-1, 0), (1, 0), (0, -1), (0, 1)]
def trap(mtx):
if not mtx:
return 0
n, m = len(mtx), len(mtx[0])
# Store perimeter cells in min heap
pq = []
seen = set()
for r in range(n):
if r in [0, n-1]:
for c in range(m):
heapq.heappush(pq, Cell(mtx[r][c], r, c))
seen.add((r, c))
else:
heapq.heappush(pq, Cell(mtx[r][0], r, 0))
seen.add((r, 0))
heapq.heappush(pq, Cell(mtx[r][m-1], r, m-1))
seen.add((r, m-1))
# Go through every seen cell stored in the heap
water = 0
popped_max = 0 # Upper bound value
while pq:
# Greedily pop off the lowest ones
popped = heapq.heappop(pq)
popped_max = max(popped_max, popped.ht)
# Examine adjacent cells
for dr, dc in SHIFT:
nr, nc = popped.r + dr, popped.c + dc
if (
0 <= nr <= n - 1 and
0 <= nc <= m - 1 and
(nr, nc) not in seen
):
# Increment water count if upper bound value
# is greater than the adjacent cell's height
water += max(popped_max - mtx[nr][nc], 0)
heapq.heappush(pq, Cell(mtx[nr][nc], nr, nc))
seen.add((nr, nc))
return water