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
83
84
85
struct Region {
  // Given m = the number of rows and n = the number of columns, (x, y) will be
  // hashed as x * n + y.
  unordered_set<int> infected;
  unordered_set<int> noninfected;
  int wallsRequired = 0;
};

class Solution {
 public:
  int containVirus(vector<vector<int>>& isInfected) {
    const int m = isInfected.size();
    const int n = isInfected[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 (isInfected[i][j] == 1 && !seen[i][j]) {
            Region region;
            // Use DFS to find all the regions (1s).
            dfs(isInfected, i, j, region, seen);
            if (!region.noninfected.empty())
              regions.push_back(region);
          }

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

      // Regions that infect the most neighbors will be sorted to the back of
      // the array.
      ranges::sort(regions, ranges::less{}, [](const Region& region) {
        return region.noninfected.size();
      });

      // Build walls around the region that infects the 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 isInfected is now contained and won't be infected anymore.
        isInfected[i][j] = 2;
      }

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

    return ans;
  }

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

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

    dfs(isInfected, i + 1, j, region, seen);
    dfs(isInfected, i - 1, j, region, seen);
    dfs(isInfected, i, j + 1, region, seen);
    dfs(isInfected, 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
78
79
80
class Region {
  // Given m = the number of rows and n = the number of columns, (x, y) will be
  // hashed as x * n + y.
  public Set<Integer> infected = new HashSet<>();
  public Set<Integer> noninfected = new HashSet<>();
  public int wallsRequired = 0;
}

class Solution {
  public int containVirus(int[][] isInfected) {
    final int m = isInfected.length;
    final int n = isInfected[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 (isInfected[i][j] == 1 && !seen[i][j]) {
            Region region = new Region();
            // Use DFS to find all the regions (1s).
            dfs(isInfected, i, j, region, seen);
            if (!region.noninfected.isEmpty())
              regions.add(region);
          }

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

      // Regions that infect the most neighbors will be sorted to the back of
      // the array.
      Collections.sort(regions, (a, b) -> a.noninfected.size() - b.noninfected.size());

      // Build walls around the region that infects the 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 isInfected is now contained and won't be infected anymore.
        isInfected[i][j] = 2;
      }

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

    return ans;
  }

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

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

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