Skip to content

749. Contain Virus 👎

  • Time: $O(m^2n^2)$
  • Space: $O(mn)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
struct Region {
  // given m = # of rows and n = # of cols, (x, y) will be hashed as x * n + y
  unordered_set<int> infected;
  unordered_set<int> noninfected;  // noninfected neighbors
  int wallsRequired = 0;
};

class Solution {
 public:
  int containVirus(vector<vector<int>>& grid) {
    const int m = grid.size();
    const int n = grid[0].size();
    int ans = 0;

    while (true) {
      vector<Region> regions;
      vector<vector<bool>> seen(m, vector<bool>(n));

      for (int i = 0; i < m; ++i)
        for (int j = 0; j < n; ++j)
          if (grid[i][j] == 1 && !seen[i][j]) {
            Region region;
            dfs(grid, i, j, region, seen);  // use DFS to find all regions (1s)
            if (!region.noninfected.empty())
              regions.push_back(region);
          }

      if (regions.empty())
        break;  // no region causes further infection

      // region which infects most neighbors is in the back
      sort(begin(regions), end(regions), [](const auto& a, const auto& b) {
        return a.noninfected.size() < b.noninfected.size();
      });

      // build walls around the region which infects most neighbors
      Region mostInfectedRegion = regions.back();
      regions.pop_back();
      ans += mostInfectedRegion.wallsRequired;

      for (const int neighbor : mostInfectedRegion.infected) {
        const int i = neighbor / n;
        const int j = neighbor % n;
        // the grid is now contained and won't be infected anymore
        grid[i][j] = 2;
      }

      // for remaining regions, expand (infect their neighbors)
      for (const Region& region : regions)
        for (const int neighbor : region.noninfected) {
          const int i = neighbor / n;
          const int j = neighbor % n;
          grid[i][j] = 1;
        }
    }

    return ans;
  }

 private:
  void dfs(const vector<vector<int>>& grid, int i, int j, Region& region,
           vector<vector<bool>>& seen) {
    if (i < 0 || i == grid.size() || j < 0 || j == grid[0].size())
      return;
    if (seen[i][j] || grid[i][j] == 2)
      return;
    if (grid[i][j] == 0) {
      region.noninfected.insert(i * grid[0].size() + j);
      ++region.wallsRequired;
      return;
    }

    // grid[i][j] == 1
    seen[i][j] = true;
    region.infected.insert(i * grid[0].size() + j);

    dfs(grid, i + 1, j, region, seen);
    dfs(grid, i - 1, j, region, seen);
    dfs(grid, i, j + 1, region, seen);
    dfs(grid, i, j - 1, region, seen);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class Region {
  // given m = # of rows and n = # of cols, (x, y) will be hashed as x * n + y
  public Set<Integer> infected = new HashSet<>();
  public Set<Integer> noninfected = new HashSet<>(); // noninfected neighbors
  public int wallsRequired = 0;
};

class Solution {
  public int containVirus(int[][] grid) {
    final int m = grid.length;
    final int n = grid[0].length;
    int ans = 0;

    while (true) {
      List<Region> regions = new ArrayList<>();
      boolean[][] seen = new boolean[m][n];

      for (int i = 0; i < m; ++i)
        for (int j = 0; j < n; ++j)
          if (grid[i][j] == 1 && !seen[i][j]) {
            Region region = new Region();
            dfs(grid, i, j, region, seen); // use DFS to find all regions (1s)
            if (!region.noninfected.isEmpty())
              regions.add(region);
          }

      if (regions.isEmpty())
        break; // no region causes further infection

      // region which infects most neighbors is in the back
      Collections.sort(regions, (a, b) -> a.noninfected.size() - b.noninfected.size());

      // build walls around the region which infects most neighbors
      Region mostInfectedRegion = regions.get(regions.size() - 1);
      regions.remove(regions.size() - 1);
      ans += mostInfectedRegion.wallsRequired;

      for (final int neighbor : mostInfectedRegion.infected) {
        final int i = neighbor / n;
        final int j = neighbor % n;
        // the grid is now contained and won't be infected anymore
        grid[i][j] = 2;
      }

      // for remaining regions, expand (infect their neighbors)
      for (final Region region : regions)
        for (final int neighbor : region.noninfected) {
          final int i = neighbor / n;
          final int j = neighbor % n;
          grid[i][j] = 1;
        }
    }

    return ans;
  }

  private void dfs(int[][] grid, int i, int j, Region region, boolean[][] seen) {
    if (i < 0 || i == grid.length || j < 0 || j == grid[0].length)
      return;
    if (seen[i][j] || grid[i][j] == 2)
      return;
    if (grid[i][j] == 0) {
      region.noninfected.add(i * grid[0].length + j);
      ++region.wallsRequired;
      return;
    }

    // grid[i][j] == 1
    seen[i][j] = true;
    region.infected.add(i * grid[0].length + j);

    dfs(grid, i + 1, j, region, seen);
    dfs(grid, i - 1, j, region, seen);
    dfs(grid, i, j + 1, region, seen);
    dfs(grid, i, j - 1, region, seen);
  }
}
Back to top