文章

class Solution {
    int[][] directions = {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
    int m;
    int n;
    int result = 0;
    int[][] g;

    boolean[][] visited;


    public int containVirus(int[][] isInfected) {
        g = isInfected;
        m = g.length;
        n = g[0].length;

        while (true) {
            int count = getCount();
            if (count == 0) {
                break;
            }
            result += count;
        }

        return result;
    }

    int getCount() {
        visited = new boolean[m][n];
        int max = -1;
        int ans = 0;
        // 连通块点集
        List<Set<Integer>> l1 = new ArrayList<>();
        // 连通块的候选感染点集
        List<Set<Integer>> l2 = new ArrayList<>();

        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (g[i][j] == 1 && visited[i][j] != true) {
                    Set<Integer> s1 = new HashSet<>();
                    Set<Integer> s2 = new HashSet<>();
                    int b = bfs(i, j, s1, s2), a = s2.size();
                    if (a > max) {
                        max = a;
                        ans = b;
                    }
                    l1.add(s1);
                    l2.add(s2);
                }
            }
        }
        for (int i = 0; i < l2.size(); i++) {
            for (int loc : l2.get(i).size() == max ? l1.get(i) : l2.get(i)) {
                int x = loc / n, y = loc % n;
                g[x][y] = l2.get(i).size() == max ? -1 : 1;
            }
        }
        return ans;

    }

    // 返回当前连通块的威胁指数
    int bfs(int i, int j, Set<Integer> s1, Set<Integer> s2) {
        int ans = 0;
        Deque<int[]> d = new ArrayDeque<>();
        visited[i][j] = true;
        d.addLast(new int[]{i, j});
        s1.add(i * n + j);
        while (!d.isEmpty()) {
            int[] info = d.pollFirst();
            int x = info[0], y = info[1];
            for (int[] di : directions) {
                int nx = x + di[0], ny = y + di[1], loc = nx * n + ny;
                if (nx < 0 || nx >= m || ny < 0 || ny >= n || visited[nx][ny]) continue;
                if (g[nx][ny] == 1) {
                    s1.add(loc);
                    visited[nx][ny] = true;
                    d.addLast(new int[]{nx, ny});
                } else if (g[nx][ny] == 0) {
                    s2.add(loc);
                    ans++;
                }
            }
        }
        return ans;
    }

}