Grokking Newsletter 203

class Solution {

  static final int NO_COLOR = 0;

  int n;
  int[][] grid;

  int[][] color;
  Map<Integer, Integer> islandSizes;

  int[][] dirs = {
    { -1, 0 }, { 0, 1 }, { 1, 0 }, { 0, -1 }
  };

  public int largestIsland(int[][] grid) {

    // setup
    this.n = grid.length;
    this.grid = grid;

    color = new int[n][n];
    islandSizes = new HashMap<>();

    int maxArea = 0;

    // find area and coloring
    int colorId = 1;

    for (int r = 0; r < n; r++) {
      for (int c = 0; c < n; c++) {

        if (grid[r][c] == 1 && color[r][c] == NO_COLOR) {

          int area = getMaxAreaAndColoring(r, c, colorId);
          islandSizes.put(colorId, area);
          colorId++;

          // record max area
          maxArea = Math.max(maxArea, area);
        }
      }
    }

    // find largest island can be created by changing one 0 cell to 1
    for (int r = 0; r < n; r++) {
      for (int c = 0; c < n; c++) {

        if (grid[r][c] == 0) {
          int area = getMaxAreaOfIslandCanMake(r, c);
          maxArea = Math.max(maxArea, area);
        }
      }
    }

    return maxArea;
  }

  int getMaxAreaAndColoring(int row, int col, int colorId) {

    Queue<int[]> queue = new LinkedList<>();
    queue.offer(new int[] { row, col });

    color[row][col] = colorId;

    int area = 0;
    while (!queue.isEmpty()) {
      int[] at = queue.poll();
      int r = at[0];
      int c = at[1];

      area++;

      for (int[] neighbor : getNeighbors(r, c)) {
        int rr = neighbor[0];
        int cc = neighbor[1];

        if (grid[rr][cc] == 1 && color[rr][cc] == NO_COLOR) {
          color[rr][cc] = colorId;

          queue.offer(new int[] { rr, cc });
        }
      }
    }

    return area;
  }

  int getMaxAreaOfIslandCanMake(int row, int col) {

    int area = 1;
    Set<Integer> seenIds = new HashSet<>();

    for (int[] neighbor : getNeighbors(row, col)) {
      int rr = neighbor[0];
      int cc = neighbor[1];

      int colorId = color[rr][cc];

      if (colorId != NO_COLOR) {
        if (!seenIds.contains(colorId)) {
          area += islandSizes.get(colorId);
        }

        seenIds.add(colorId);
      }
    }

    return area;
  }

  List<int[]> getNeighbors(int row, int col) {

    List<int[]> ans = new ArrayList<>();

    for (int[] dir : dirs) {
      int rr = row + dir[0];
      int cc = col + dir[1];

      if (0 <= rr && rr < n && 0 <= cc && cc < n) {
        ans.add(new int[] { rr, cc });
      }
    }

    return ans;
  }
}