题目描述

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.

Example:

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.

时间复杂度 O(mn)

Java

class Solution {
    class Cell implements Comparable<Cell> {
        private int row;
        private int column;
        private int height;

        public Cell(int row, int column, int height) {
            this.row = row;
            this.column = column;
            this.height = height;
        }

        @Override
        public int compareTo(Cell o) {
            return this.height - o.height;
        }
    }

    int[] dx = {0, 1, 0, -1};
    int[] dy = {1, 0, -1, 0};
    boolean[][] visit = null;

    public int trapRainWater(int[][] heightMap) {
        if (heightMap.length == 0 || heightMap[0].length == 0) {
            return 0;
        }
        int m = heightMap.length;
        int n = heightMap[0].length;
        visit = new boolean[m][n];
        PriorityQueue<Cell> pq = new PriorityQueue<>();

        // 第一行，最后一行
        for (int j = 0; j < n; j++) {
            pq.offer(new Cell(0, j, heightMap[0][j]));
            pq.offer(new Cell(m - 1, j, heightMap[m - 1][j]));
            visit[0][j] = true;
            visit[m - 1][j] = true;
        }

        // 第一列，最后一列
        for (int i = 0; i < m; i++) {
            pq.offer(new Cell(i, 0, heightMap[i][0]));
            pq.offer(new Cell(i, n - 1, heightMap[i][n - 1]));
            visit[i][0] = true;
            visit[i][n - 1] = true;
        }

        int res = 0;
        while (!pq.isEmpty()) {
            Cell curr = pq.poll();
            int h = curr.height;
            int r = curr.row;
            int c = curr.column;

            for (int i = 0; i < 4; i++) {
                int x = r + dx[i];
                int y = c + dy[i];
                if (x > 0 && y > 0 && x < m - 1 && y < n - 1 && !visit[x][y]) {
                    visit[x][y] = true;
                    res += Math.max(0, h - heightMap[x][y]);
                    pq.offer(new Cell(x, y, Math.max(h, heightMap[x][y])));
                }
            }
        }
        return res;
    }
}