题目地址：

leetcode.com/problems/mi…

给定一个n个节点的无向图，以邻接矩阵给出。每个节点代表一个网络中的节点。再给定一个数组 A，A[i]表示节点 i被感染了（下文称 A里的点是”感染源头“）。当一个节点被感染了，它所在的连通块的所有节点都会被感染。现在允许选取一个“感染源头”，将其删掉（删掉之后，这个感染源和其所有邻接边都会被删除）。问将哪个节点删掉可以使得剩余节点的总的被感染节点数量最小。如果有多个答案，则返回编号最小的节点。

我们主要是要考虑删掉某个感染源后，能“拯救”多少个节点。我们可以先不考虑感染源，然后用并查集统计连通块。接着，我们用哈希表统计一下每个连通块被多少个感染源感染了，略过没有被感染的连通块，如果其被多于 1个感染源感染了，那么无论删掉哪个感染源该连通块都无法被拯救，这样的连通块就不考虑了；如果其恰好是被 1个感染源感染了，那么将这个感染源删掉能拯救的节点数，就是这个连通块元素个数。我们枚举每个感染源，然后看删去它能拯救多少节点，找到能拯救最多节点的那个感染源即可。如果答案不唯一则找编号最小的那个。

import java.util.*;

public class Solution {
    
    class UnionFind {
        int[] parent, size;
        
        public UnionFind(int size) {
            parent = new int[size];
            for (int i = 0; i < size; i++) {
                parent[i] = i;
            }
            this.size = new int[size];
            Arrays.fill(this.size, 1);
        }
        
        public int find(int x) {
            if (parent[x] != x) {
                parent[x] = find(parent[x]);
            }
            return parent[x];
        }
        
        public void union(int x, int y) {
            int px = find(x), py = find(y);
            if (px == py) {
                return;
            }
            
            parent[px] = py;
            size[py] += size[px];
        }
    }
    
    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;
        // 记录一下哪些点是感染源
        boolean[] dirty = new boolean[n];
        for (int x : initial) {
            dirty[x] = true;
        }
        
        // 先不考虑感染源，求一下连通块
        UnionFind uf = new UnionFind(n);
        for (int i = 0; i < n; i++) {
            if (!dirty[i]) {
                for (int j = i + 1; j < n; j++) {
                    if (!dirty[j] && graph[i][j] == 1) {
                        uf.union(i, j);
                    }
                }
            }
        }
        
        // 求一下每个连通块被哪些感染源感染了，key存连通块的树根，value存能感染到该连通块的感染源编号
        Map<Integer, Set<Integer>> infectorMap = new HashMap<>();
        for (int x : initial) {
            for (int y = 0; y < graph[x].length; y++) {
            	// 如果通过x能感染到y，则用哈希表存一下
                if (!dirty[y] && graph[x][y] == 1) {
                    int py = uf.find(y);
                    infectorMap.putIfAbsent(py, new HashSet<>());
                    infectorMap.get(py).add(x);
                }
            }
        }
        
        // 求一下每个感染源如果将其删掉，能拯救多少个节点。
        // 注意这里要略过被多个感染源感染的连通块，这些连通块是无法被拯救的
        Map<Integer, Integer> saveCount = new HashMap<>();
        for (Map.Entry<Integer, Set<Integer>> entry : infectorMap.entrySet()) {
            int parent = entry.getKey();
            Set<Integer> infectors = entry.getValue();
            // 如果当前连通块被多个感染源感染了，那就略过
            if (infectors.size() > 1) {
                continue;
            }
            
            // 累加一下这个感染源如果删掉能拯救多少个节点
            int infector = infectors.iterator().next();
            saveCount.put(infector, saveCount.getOrDefault(infector, 0) + uf.size[parent]);
        }
        
        // 如果每个连通块都被多于1个感染源感染了，那就删掉编号最小的感染源
        if (saveCount.isEmpty()) {
            int res = n;
            for (int i : initial) {
                res = Math.min(res, i);
            }
            
            return res;
        }
        
        // 否则去找一下被拯救节点最多的那个感染源
        int res = -1, maxSave = 0;
        for (Map.Entry<Integer, Integer> entry : saveCount.entrySet()) {
            int infector = entry.getKey(), saved = entry.getValue();
            if (saved > maxSave) {
                maxSave = saved;
                res = infector;
            } else if (saved == maxSave) {
                res = Math.min(res, infector);
            }
        }
        
        return res;
    }
}

【Leetcode】928. Minimize Malware Spread II（代码）

题目地址：