leetcode-go/leetcode/hard/407.trapping-rain-water-ii.java

import java.util.PriorityQueue;

class Solution {
    class Grid implements Comparable<Grid> {
        int val;
        int y;
        int x;

        Grid(int val, int y, int x) {
            this.val = val;
            this.y = y;
            this.x = x;
        }

        @Override
        public int compareTo(Grid that) {
            return this.val - that.val; // Don't forget the "public"
        }
    }

    int[] dy = new int[]{0, 1, 0, -1};
    int[] dx = new int[]{1, 0, -1, 0};

    public int trapRainWater(int[][] heightMap) {
        int m, n;
        if ((m = heightMap.length) == 0) return 0;
        if ((n = heightMap[0].length) == 0) return 0;
        PriorityQueue<Grid> pq = new PriorityQueue<>();
        boolean[][] used = new boolean[m][n];
        for (int i = 0; i < m; i++) {
            pq.offer(new Grid(heightMap[i][0], i, 0));
            pq.offer(new Grid(heightMap[i][n - 1], i, n - 1));
            used[i][0] = true;
            used[i][n - 1] = true;
        }
        for (int j = 1; j < n - 1; j++) {
            pq.offer(new Grid(heightMap[0][j], 0, j));
            pq.offer(new Grid(heightMap[m - 1][j], m - 1, j));
            used[0][j] = true;
            used[m - 1][j] = true;
        }
        int water = 0, height = 0;
        while (!pq.isEmpty()) {
            Grid curr = pq.poll();
            if (height < curr.val) height = curr.val;
            else water += height - curr.val;
            for (int i = 0; i < 4; i++) {
                int ny = curr.y + dy[i], nx = curr.x + dx[i];
                if (ny < 0 || m <= ny || nx < 0 || n <= nx || used[ny][nx]) continue;
                pq.offer(new Grid(heightMap[ny][nx], ny, nx));
                used[ny][nx] = true;
            }
        }
        return water;
    }
}