Java solution using PriorityQueue


  • 106
    Y

    Source code from:
    https://github.com/shawnfan/LintCode/blob/master/Java/Trapping Rain Water II.java

    
    public class Solution {
    
        public class Cell {
            int row;
            int col;
            int height;
            public Cell(int row, int col, int height) {
                this.row = row;
                this.col = col;
                this.height = height;
            }
        }
    
        public int trapRainWater(int[][] heights) {
            if (heights == null || heights.length == 0 || heights[0].length == 0)
                return 0;
    
            PriorityQueue<Cell> queue = new PriorityQueue<>(1, new Comparator<Cell>(){
                public int compare(Cell a, Cell b) {
                    return a.height - b.height;
                }
            });
            
            int m = heights.length;
            int n = heights[0].length;
            boolean[][] visited = new boolean[m][n];
    
            // Initially, add all the Cells which are on borders to the queue.
            for (int i = 0; i < m; i++) {
                visited[i][0] = true;
                visited[i][n - 1] = true;
                queue.offer(new Cell(i, 0, heights[i][0]));
                queue.offer(new Cell(i, n - 1, heights[i][n - 1]));
            }
    
            for (int i = 0; i < n; i++) {
                visited[0][i] = true;
                visited[m - 1][i] = true;
                queue.offer(new Cell(0, i, heights[0][i]));
                queue.offer(new Cell(m - 1, i, heights[m - 1][i]));
            }
    
            // from the borders, pick the shortest cell visited and check its neighbors:
            // if the neighbor is shorter, collect the water it can trap and update its height as its height plus the water trapped
           // add all its neighbors to the queue.
            int[][] dirs = new int[][]{{-1, 0}, {1, 0}, {0, -1}, {0, 1}};
            int res = 0;
            while (!queue.isEmpty()) {
                Cell cell = queue.poll();
                for (int[] dir : dirs) {
                    int row = cell.row + dir[0];
                    int col = cell.col + dir[1];
                    if (row >= 0 && row < m && col >= 0 && col < n && !visited[row][col]) {
                        visited[row][col] = true;
                        res += Math.max(0, cell.height - heights[row][col]);
                        queue.offer(new Cell(row, col, Math.max(heights[row][col], cell.height)));
                    }
                }
            }
            
            return res;
        }
    }
    

  • 19
    D

    @yuhaowang001 Your idea is interesting because it maintains a boundary.

    This solution is also correct because of the invariant: the boundary is always in the queue.

    One tricky case that this solution handles very well is :

    9 9 9 9 9 9 8 9 9 9 9
    9 0 0 0 0 0 1 0 0 0 9
    9 0 0 0 0 0 0 0 0 0 9
    9 0 0 0 0 0 0 0 0 0 9
    9 9 9 9 9 9 9 9 9 9 9

    After you process 8, the downward 1 will be replaced as 8, instead of 1 as height.


  • 0
    D

    @yuhaowang001 Just practice your idea :)

    public class Solution {
        public int trapRainWater(int[][] heightMap) {
            if (heightMap == null || heightMap.length == 0 || heightMap[0].length == 0) {
                return 0;
            }
            int rows = heightMap.length;
            int columns = heightMap[0].length;
            Queue<Cell> minheap = new PriorityQueue<>(rows * columns, new Comparator<Cell>() {
                @Override
                public int compare(Cell cell1, Cell cell2) {
                    return cell1.h - cell2.h;
                }
            });
            boolean[][] visited = new boolean[rows][columns];
            for (int c = 0; c < columns - 1; c++) {
                minheap.offer(new Cell(0, c, heightMap[0][c]));
                visited[0][c] = true;
            }
            for (int r = 0; r < rows - 1; r++) {
                minheap.offer(new Cell(r, columns - 1, heightMap[r][columns - 1]));
                visited[r][columns - 1] = true;
            }
            for (int c = columns - 1; c > 0; c--) {
                minheap.offer(new Cell(rows - 1, c, heightMap[rows - 1][c]));
                visited[rows - 1][c] = true;
            }
            for (int r = rows - 1; r > 0; r--) {
                minheap.offer(new Cell(r, 0, heightMap[r][0]));
                visited[r][0] = true;
            }
            int water = 0;
            int[] xplus = new int[] {-1, 1, 0, 0};
            int[] yplus = new int[] {0, 0, -1, 1};
            while (!minheap.isEmpty()) {
                Cell cell = minheap.poll();
                for (int i = 0; i < 4; i++) {
                    int newx = cell.x + xplus[i];
                    int newy = cell.y + yplus[i];
                    if (newx < 0 || newx >= rows || newy < 0 || newy >= columns || visited[newx][newy]) {
                        continue;
                    }
                    visited[newx][newy] = true;
                    water += Math.max(0, cell.h - heightMap[newx][newy]);
                    minheap.offer(new Cell(newx, newy, Math.max(cell.h, heightMap[newx][newy])));
                }
            }
            return water;
        }
        
        class Cell {
            int x;
            int y;
            int h;
            public Cell(int x, int y, int h) {
                this.x = x;
                this.y = y;
                this.h = h;
            }
        }
    }
    

  • 21

    Brilliant solution!
    Two little improvements about your code:

    1. The first line could be if(heightMap == null || heightMap.length <= 2 || heightMap[0].length <= 2) return 0; because 1-D array and 2-D matrix whose length is 2 are not valid for holding water.
    2. You could write compareTo method in Cell, like:
    class Cell implements Comparable<Cell> {
            public int row, col, height;
    
            public Cell(int row, int col, int height) {
                this.row = row;
                this.col = col;
                this.height = height;
            }
    
            public int compareTo(Cell o) {
                return this.height - o.height;
            }
        }
    

    In this way, you don't need to write compare method for PriorityQueue.


  • 0
    D

    @xietao0221 I am still trying to figure out the relation between Comparable/Comparator interface and equals() method. It looks like Comparable/Comparator interface has to be consistent with equals(), that is to say, when Comparator says equal, equals() has to be true. Not sure it is an "if and only if" relation. But anyway, it is something that we might need to pay attention to.


  • 2

    @dachuan.huang in my opinion, Comparable/Comparator is used to do the sorting, for example, when you call Collections.sort() or Arrays.sort() or PriorityQueue related functions, you need to compose Comparable/Comparator. But equals is used to do hashing related functions, for example, when you create a class which is need to be hashed, you have to override equals and hashCode method. Correct me if I am wrong.


  • 3

    @xietao0221 even better if you change 1 to 2.


  • 1

    @mzchen Yes, you are right! I'll modify my post.


  • 0

    @yuhaowang001 Nice solution! Would you mind sharing what characteristic of the problem suggested you to use a pq? In another word, how did you come up with the solution? Thanks in advance.


  • 1
    I

    Almost same idea, but used DFS for the cells that are shorter or equal, because you know they are going to be switched to the root of the queue, why bother putting them back wasting logN time.

    150ms, the lambda function took time, if replaced with old style it reduced to 30ms, right into the bell shape distribution.

    final static int[][] dirs = {{0, 1}, {1, 0}, {0, -1}, {-1, 0}};
    public int trapRainWater(int[][] heightMap) {
        if (heightMap.length == 0) return 0;
        Queue<int[]> pq = new PriorityQueue<>((a, b) -> heightMap[a[0]][a[1]] - heightMap[b[0]][b[1]]);
        int m = heightMap.length, n = heightMap[0].length, ans = 0, x = 0, y = 0, len = n;
        boolean[][] visited = new boolean[m][n];
        for (int[] dir : dirs) {
            for (int i = 0;; ++i) {
                visited[x][y] = true;
                pq.offer(new int[]{x, y});
                if (i == len - 1) break;
                x += dir[0];
                y += dir[1];
            }
            len ^= m ^ n;
        }
        while (!pq.isEmpty()) {
            int[] p = pq.poll();
            ans += visit(p, visited, heightMap, heightMap[p[0]][p[1]], pq);
        }
        return ans;
    }
    int visit(int[] pos, boolean[][] visited, int[][] hMap, int h, Queue<int[]> pq) {
        int m = hMap.length, n = hMap[0].length, ans = 0, x = pos[0], y = pos[1];
        visited[x][y] = true;
        if (h < hMap[x][y]) {
            pq.offer(pos);
        } else {
            ans += h - hMap[x][y];
            for (int[] dir : dirs) {
                int[] p = new int[]{x + dir[0], y + dir[1]};
                if (p[0] < 0 || p[0] == m || p[1] < 0 || p[1] == n || visited[p[0]][p[1]]) continue;
                ans += visit(p, visited, hMap, h, pq);
            }
        }
        return ans;
    }
    

  • 0
    S

    What's the time and space complexity for this solution?


  • 1
    H

    excellent idea! thanks for sharing your implementation!


  • 0
    S

    Genius! Brilliant!


  • 0

    Any explanations for the code? I do not think the correctness of the algorithm is obvious.


  • 7
    Y

    @shaokunzou The correctness of this algorithm is based on the fact that each time we only examine the boarder point with the smallest height and we traverse from outside boarder to insider point. So each time you encounter with a point in the graph with a smaller height, you are guaranteed to be able to hold at least cell.height - heights[row][col] water in this point (remember that current cell we used are the one with the smallest height among the boarders).


  • 4
    Y

    @sherlywang The time complexity consists of 2 parts: the heapify process and the while loop.
    We put m + n elements into the heap during the heapify process, so it's O(m + n) run time.
    During the while loop, every cell is put into and take out of the heap at most once, and we are doing so in a BFS style, meaning that there is m + n elements in the heap at the worst case. So it is O(m * n * log(m + n)) in the worst case.
    So the run time complexity would be O(m * n * log(m + n)).
    The space complexity is obviously O(m * n) because of the visited array.


  • 2
    B

    When you initialize the priority queue,
    for (int i = 0; i < m; i++)
    for (int i = 0; i < n; i++)
    are sued. It seems like you add the four corner cells into queue twice?
    The second for loop should be
    for (int i = 1; i < n-1; i++) {...}


  • 0
    Y

    @BirdFrank Yes, you are right. But it still works putting these cells into the queue twice though it causes extra workload. When the same cell is polled the second time, its visited neighbors would not be visited one more time.


  • -3
    C

    Another very interesting case this code can handler very well is:
    9 9 9 9 9
    9 2 1 2 9
    9 2 8 2 9
    9 2 3 2 9
    9 9 9 9 9

    Where no water can be trapped


  • 0
    N

    @coderliang : Water will be trapped in the example given by you. Why you think it wont be trapped?


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.