Reference: LeetCode
Difficulty: Hard

Problem

Given an m x n matrix of positive integers representing the height of each unit cell in a 2D elevation map, compute the volume of water it is able to trap after raining.

Note: Both m and n are less than 110. The height of each unit cell is greater than 0 and is less than 20,000.

Example: 

1
2
3
4
5
6
Given the following 3x6 height map: [
  [1,4,3,1,3,2],
  [3,2,1,3,2,4],
  [2,3,3,2,3,1]
]
Return 4.

Analysis

Priority Queue

Explanation: Animation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
class Cell {
  int i; int j; int height;
  Cell(int _i, int _j, int _height) { i = _i; j = _j; height = _height; }
}

//       n
// [1,4,3,1,3,2],
// [3,2,1,3,2,4],   m
// [2,3,3,2,3,1]
public int trapRainWater(int[][] heightMap) {
  if (heightMap == null || heightMap.length == 0 || heightMap[0].length == 0) {
    return 0;
  }
  int m = heightMap.length;
  int n = heightMap[0].length;
  boolean[][] visited = new boolean[m][n]; // false. If added to pq, become true
  PriorityQueue<Cell> pq = new PriorityQueue<>((c1 , c2) -> (c1.height - c2.height)); // min pq
  // init pq with border elements
  for (int col = 0; col < n; ++col) {
    visit(0, col, pq, visited, heightMap); // first row
    visit(m - 1, col, pq, visited, heightMap); // last row
  }
  for (int row = 1; row < m - 1; ++row) { // skip the first and the last row
    visit(row, 0, pq, visited, heightMap);
    visit(row, n - 1, pq, visited, heightMap);
  }
  // remove and add
  int max = Integer.MIN_VALUE;
  int water = 0;
  while (pq.size() > 0) {
    Cell c = pq.remove();
    int ht = c.height;
    if (ht <= max) water += max - ht;
    else          max = ht; // update max
    // visit its unvisited neighbors
    int row = c.i, col = c.j;
    if (row - 1 >= 0 && !visited[row - 1][col]) { // up
      visit(row - 1, col, pq, visited, heightMap);
    }
    if (row + 1 <= m - 1 && !visited[row + 1][col]) { // down
      visit(row + 1, col, pq, visited, heightMap);
    } 
    if (col - 1 >= 0 && !visited[row][col - 1]) { // left
      visit(row, col - 1, pq, visited, heightMap);
    }
    if (col + 1 <= n - 1 && !visited[row][col + 1]) { // right
      visit(row, col + 1, pq, visited, heightMap);
    }
  }
  return water;
}

// add to pq, and set visited true
private void visit(int row, int col, PriorityQueue<Cell> pq, boolean[][] visited, int[][] heightMap) {
  pq.add(new Cell(row, col, heightMap[row][col]));
  visited[row][col] = true;
}

Time: $O(MN\log{(MN)})$
Space: $O(MN)$