924. Minimize Malware Spread 

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 Solution {
    private static class UnionFind {

        private final int[] parent;
        private final int[] size;

        public UnionFind(int n) {
            this.parent = new int[n];
            this.size = new int[n];
            for (int i = 0; i < parent.length; i++) {
                parent[i] = i;
                size[i] = 1;
            }
        }

        public void union(int x, int y) {
            int xRoot = getRoot(x), yRoot = getRoot(y);
            if (xRoot == yRoot) {
                return;
            }

            if (size[xRoot] <= size[yRoot]) {
                parent[xRoot] = yRoot;
                size[yRoot] += size[xRoot];
            } else {
                parent[yRoot] = xRoot;
                size[xRoot] += size[yRoot];
            }
        }

        public int getRoot(int x) {
            if (parent[x] != x) {
                parent[x] = getRoot(parent[x]);
            }

            return parent[x];
        }

        public int getSize(int root) {
            return size[root];
        }

    }

    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;

        UnionFind unionFind = new UnionFind(n);
        for (int i = 0; i < n; i++) {
            for (int j = i + 1; j < n; j++) {
                if (graph[i][j] == 1) {
                    unionFind.union(i, j);
                }
            }
        }

        int minX = n;
        int minInitial = n;
        int maxEliminations = 0;
        int[] rootToInitialSizeMap = new int[n];
        for (int x : initial) {
            rootToInitialSizeMap[unionFind.getRoot(x)]++;
            minInitial = Math.min(minInitial, x);
        }

        for (int x : initial) { // 计算移除当前节点可以减少感染的节点数
            int root = unionFind.getRoot(x);
            if (rootToInitialSizeMap[root] == 1) {
                int eliminations = unionFind.getSize(root);
                // 注意对 minX 的更新逻辑是满足明确的条件才能更新，不能简单使用 Math.min 进行更新
                if (eliminations > maxEliminations || (eliminations == maxEliminations && x < minX)) {
                    maxEliminations = eliminations;
                    minX = x;
                }
            }
        }

        return minX == n ? minInitial : minX;
    }
}

References

924. Minimize Malware Spread