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; } }
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]; PriorityQueue<Cell> pq = new PriorityQueue<>((c1 , c2) > (c1.height  c2.height)); for (int col = 0; col < n; ++col) { visit(0, col, pq, visited, heightMap); visit(m  1, col, pq, visited, heightMap); } for (int row = 1; row < m  1; ++row) { visit(row, 0, pq, visited, heightMap); visit(row, n  1, pq, visited, heightMap); } 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; int row = c.i, col = c.j; if (row  1 >= 0 && !visited[row  1][col]) { visit(row  1, col, pq, visited, heightMap); } if (row + 1 <= m  1 && !visited[row + 1][col]) { visit(row + 1, col, pq, visited, heightMap); } if (col  1 >= 0 && !visited[row][col  1]) { visit(row, col  1, pq, visited, heightMap); } if (col + 1 <= n  1 && !visited[row][col + 1]) { visit(row, col + 1, pq, visited, heightMap); } } return water; }
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; }
