Contain Virus LeetCode solution in Java

1
class Solution {
2

3
  private static final int[][] DIR = new int[][] {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
4

5
  public int containVirus(int[][] isInfected) {
6
    int m = isInfected.length, n = isInfected[0].length;
7
    int ans = 0;
8

9
    while (true) {
10
      // infected regions, sorted desc according to the number of nearby
11
      // uninfected nodes
12
      PriorityQueue<Region> pq = new PriorityQueue<Region>();
13
      // already visited cells
14
      boolean[][] visited = new boolean[m][n];
15

16
      // find regions
17
      for (int i = 0; i < m; i++) {
18
        for (int j = 0; j < n; j++) {
19

20
          // if current cell is infected, and it's not visited
21
          if (isInfected[i][j] != 1 || visited[i][j]) continue;
22

23
          // we found a new region, dfs to find all the infected
24
          // and uninfected cells in the current region
25
          Region reg = new Region();
26
          dfs(i, j, reg, isInfected, visited, new boolean[m][n], m, n);
27

28
          // if there are some uninfected nodes in this region,
29
          // we can contain it, so add it to priority queue
30
          if (reg.uninfected.size() != 0) pq.offer(reg);
31
        }
32
      }
33

34
      // if there are no regions to contain, break
35
      if (pq.isEmpty()) break;
36

37
      // Contain region with most uninfected nodes
38
      Region containReg = pq.poll();
39
      ans += containReg.wallsRequired;
40

41
      // use (2) to mark a cell as contained
42
      for (int[] cell : containReg.infected) isInfected[cell[0]][cell[1]] = 2;
43

44
      // Spread infection to uninfected nodes in other regions
45
      while (!pq.isEmpty()) {
46
        Region spreadReg = pq.poll();
47

48
        for (int[] cell : spreadReg.uninfected) isInfected[cell[0]][cell[1]] = 1;
49
      }
50
    }
51
    return ans;
52
  }
53

54
  private void dfs(
55
      int i,
56
      int j,
57
      Region reg,
58
      int[][] grid,
59
      boolean[][] visited,
60
      boolean[][] uninfectedVis,
61
      int m,
62
      int n) {
63
    visited[i][j] = true;
64
    reg.addInfected(i, j);
65

66
    for (int[] dir : DIR) {
67
      int di = i + dir[0];
68
      int dj = j + dir[1];
69

70
      // continue, if out of bounds OR contained OR already visited
71
      if (di < 0 || dj < 0 || di == m || dj == n || grid[di][dj] == 2 || visited[di][dj]) continue;
72

73
      // if neighbour node is not infected
74
      if (grid[di][dj] == 0) {
75
        // a wall will require to stop the spread from cell (i,j) to (di, dj)
76
        reg.wallsRequired++;
77

78
        // if this uninfected node is not already visited for current region
79
        if (!uninfectedVis[di][dj]) {
80
          uninfectedVis[di][dj] = true;
81
          reg.addUninfected(di, dj);
82
        }
83
      } else dfs(di, dj, reg, grid, visited, uninfectedVis, m, n);
84
    }
85
  }
86
}
87

88
class Region implements Comparable<Region> {
89
  public List<int[]> infected;
90
  public List<int[]> uninfected;
91
  public int wallsRequired;
92

93
  public Region() {
94
    infected = new ArrayList();
95
    uninfected = new ArrayList();
96
  }
97

98
  public void addInfected(int row, int col) {
99
    infected.add(new int[] {row, col});
100
  }
101

102
  public void addUninfected(int row, int col) {
103
    uninfected.add(new int[] {row, col});
104
  }
105

106
  @Override
107
  public int compareTo(Region r2) {
108
    return Integer.compare(r2.uninfected.size(), uninfected.size());
109
  }
110
}