【Leetcode】924. Minimize Malware Spread

题目地址：

https://leetcode.com/problems/minimize-malware-spread/

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

这个问题可以和https://blog.csdn.net/qq_46105170/article/details/113489000一起看，两者解答可以互相借鉴。DFS做法参考https://blog.csdn.net/qq_46105170/article/details/113489029。下面介绍并查集做法。

我们主要是要考虑“清洁”某个感染源后，能“拯救”多少个节点。如果对于连通块$C_1,...,C_k$，它们的”感染源“数都是$1$，那显然要找到那个最大的连通块，清洁那个最大的连通块的那个”感染源“的话能”拯救“最多的节点。排除掉没被感染的连通块后，对于”感染源“多于$1$的连通块，这些连通块无论清洁哪个”感染源“都只能”拯救“$1$个节点（就是被清洁的节点自己），所以按照要求，就返回$A$里的最小值即可。具体算法可以这样叙述：
1、先用并查集将所有被感染的连通块找出来；
2、算一下每个被感染的连通块里的感染源个数，用哈希表记录下来；
3、遍历哈希表，如果当前被感染的连通块里感染源个数多于$1$则略过，否则这个连通块的size就是将它里面的感染源“清洁”之后能拯救的节点数。我们找到能拯救最多节点的感染源编号，如果有多个解则找到编号最小者，将之返回。
4、如果每个被感染的连通块里的感染源数都多于$1$，那“清洁”任意一个感染源都只能拯救$1$个节点，那就返回感染源里的编号最小者即可。

代码如下：

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

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;
        UnionFind uf = new UnionFind(n);
        
        // 矩阵是对称的，所以只需要遍历右上半部分就可以了
        for (int i = 0; i < n; i++) {
            for (int j = i + 1; j < n; j++) {
                if (graph[i][j] == 1) {
                    uf.union(i, j);
                }
            }
        }
        
        // 用map存一下每个被感染的连通块里，”感染源“的个数
        Map<Integer, Integer> map = new HashMap<>();
        for (int x : initial) {
            int px = uf.find(x);
            map.put(px, map.getOrDefault(px, 0) + 1);
        }
        
        // maxN存被感染的最大连通块的size
        int maxN = 0, res = -1;
        for (int x : initial) {
            int px = uf.find(x);
            // 如果找到了感染源数是1的连通块，那么就挑一个size最大的连通块
            if (map.get(px) == 1) {
                if (uf.size[px] > maxN) {
                    maxN = uf.size[px];
                    res = x;
                } else if (uf.size[px] == maxN) {
                	// 如果有多个感染源数是1的连通块size一样，
                	// 则按照要求要返回编号最小的感染源头节点
                    res = Math.min(res, x);
                }
            }
        }
        
        // 如果每个连通块的感染源都多于1，那就清洁initial里编号最小的节点
        if (res == -1) {
            res = n;
            for (int x : initial) {
                res = Math.min(res, x);
            }
        }
        
        return res;
    }
}

时间复杂度$O(n^2{\log }^{\ast }n)$，空间$O(n)$。