928. Minimize Malware Spread II 

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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
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;

        boolean[] initialSet = new boolean[n];
        for (int x : initial) {
            initialSet[x] = true;
        }

        UnionFind unionFind = new UnionFind(n);
        // 仅将非初始感染节点构建为并查集
        for (int i = 0; i < n; i++) {
            if (initialSet[i]) {
                continue;
            }
            for (int j = i + 1; j < n; j++) {
                if (initialSet[j]) {
                    continue;
                }
                if (graph[i][j] == 1) {
                    unionFind.union(i, j);
                }
            }
        }

        int[] rootToInitialCountMap = new int[n];
        Map<Integer, Set<Integer>> initialToRootSetMap = new HashMap<>();

        for (int x : initial) {
            for (int i = 0; i < n; i++) {
                if (graph[x][i] == 1 && !initialSet[i]) {
                    int root = unionFind.getRoot(i);
                    initialToRootSetMap.computeIfAbsent(x, k -> new HashSet<>()).add(root);
                }
            }

            for (int root : initialToRootSetMap.getOrDefault(x, Collections.emptySet())) {
                rootToInitialCountMap[root]++; // 注意该计数需要在上方 for 循环完成后进行，避免单个 root 被多次统计，即单个初始节点的相邻节点均属于同一个连通分量的情况
            }
        }

        int minX = n;
        int minInitial = n;
        int maxEliminations = 0;

        for (int x : initial) {
            minInitial = Math.min(minInitial, x);

            int eliminations = 0;
            for (int root : initialToRootSetMap.getOrDefault(x, Collections.emptySet())) {
                if (rootToInitialCountMap[root] == 1) {
                    eliminations += unionFind.getSize(root);
                }
            }

            if (eliminations > maxEliminations || (eliminations == maxEliminations && x < minX)) {
                maxEliminations = eliminations;
                minX = x;
            }
        }

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

References

928. Minimize Malware Spread II