LeetCode/java/trapping-rain-water-ii.java

class Solution {
  private int mapWidth = -1;
  private int mapHeight = -1;
  private int[][] heightMap = null;

  private PriorityQueue<Coord> q = null;
  private boolean[][] visited = null;

  private boolean isValid(int y, int x) {
    return y >= 0 && y < mapHeight && x >= 0 && x < mapWidth;
  }

  record Coord(int height, int y, int x) implements Comparable<Coord> {

    public Coord add(Solution s, Coord other) {
      var nextY = y + other.y;
      var nextX = x + other.x;
      if (!s.isValid(nextY, nextX)) {
        return null;
      }

      return new Coord(Math.max(height, s.heightMap[nextY][nextX]), nextY, nextX);
    }

    @Override
    public int compareTo(Coord other) {
      return Integer.compare(this.height, other.height);
    }
  }

  private static final Coord[] DIRECTIONS =
      new Coord[] {
        new Coord(0, 0, -1), new Coord(0, 0, 1), new Coord(0, -1, 0), new Coord(0, 1, 0)
      };

  private void initQueue(int n, int y0, int x0, int dy, int dx) {
    for (int i = 0; i < n; ++i) {
      var y = y0 + i * dy;
      var x = x0 + i * dx;

      q.offer(new Coord(heightMap[y][x], y, x));
      visited[y][x] = true;
    }
  }

  public int trapRainWater(int[][] heightMap) {
    mapWidth = heightMap[0].length;
    mapHeight = heightMap.length;
    this.heightMap = heightMap;

    visited = new boolean[mapHeight][mapWidth];
    q = new PriorityQueue<>();

    // initialize outer columns
    initQueue(mapHeight, 0, 0, 1, 0);
    initQueue(mapHeight, 0, mapWidth - 1, 1, 0);

    // initialize outer rows
    initQueue(mapWidth, 0, 0, 0, 1);
    initQueue(mapWidth, mapHeight - 1, 0, 0, 1);

    var totalVolume = 0;
    while (!q.isEmpty()) {
      var cell = q.poll();

      for (var dir : DIRECTIONS) {
        var neighbor = cell.add(this, dir);
        if (neighbor == null || visited[neighbor.y()][neighbor.x()]) {
          continue;
        }

        var neighborHeight = heightMap[neighbor.y()][neighbor.x()];
        if (neighborHeight < cell.height()) {
          totalVolume += cell.height() - neighborHeight;
        }

        q.offer(neighbor);
        visited[neighbor.y()][neighbor.x()] = true;
      }
    }

    return totalVolume;
  }
}