928. 尽量减少恶意软件的传播 II

原题链接

我们先根据题意构建图G。然后我们可以得出如果一个普通结点A连接着两个感染源结点，那么其无论如何都会被感染；因此我们要找寻那些只与一个感染源结点相连接的普通节点；

然后我们在图中把感染源结点擦除，可以得到几个只由普通结点构成的图，我们把这几张图看作一个集合（即使用并查集）；接着我们遍历每个集合，得到该集合是否只与一个感染源结点相连接，并累加。

Code

class Solution {
public:
    static const int N=305;
　　//father负责并查集的实现
　　//virus负责对结点是否为感染源结点进行标记
　　//Size负责统计集合内元素个数
　　//cnt负责统计每个集合连接的感染源结点下标（-1，表示未找到，-2表示有两个以上感染源结点连接）
　　//sum负责统计每个感染源结点去除后可恢复多少普通结点
    int father[N],virus[N],n,Size[N],cnt[N],sum[N];
    void build(vector<int>& initial){  //初始化
        for (int i=0;i<n;i++){
            father[i]=i;
            Size[i]=1;
            cnt[i]=-1;
        }
        for (auto &i : initial){
            virus[i]=1;
        }
    }
    int find(int x){
        if (x!=father[x])
            father[x]=find(father[x]);
        return father[x];
    }
    void Union(int x,int y){
        int fx=find(x);
        int fy=find(y);
        if (fx!=fy){
            father[fx]=fy;
            Size[fy]+=Size[fx];
        }
    }
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
        n=graph.size();
        int n2=initial.size();
        build(initial);
        for (int i=0;i<n;i++){
            for (int j=0;j<n;j++){
                if (j!=i && virus[i]==0 && graph[i][j]==1 && virus[j]==0) Union(i,j);  //对普通结点进行合并
            }
        }
        for (int i=0;i<n2;i++){
            for (int j=0;j<n;j++){
                if (j!=initial[i] && virus[j]==0 && graph[initial[i]][j]==1){   //找寻每个集合所连接的感染源结点
                    int fj=find(j);
                    if (cnt[fj]==-1) cnt[fj]=initial[i];
                    else if (cnt[fj]!=initial[i]) cnt[fj]=-2;
                } 
            }
        }
        sort(initial.begin(),initial.end());
        for (int i=0;i<n;i++){
            if (father[i]==i && virus[i]==0 && cnt[i]>=0){  //统计每个感染源结点去除后可恢复多少普通结点
                sum[cnt[i]]+=Size[i];
            }
        }
        int Max=0,cnt_i=initial[0];
        for (int i=0;i<n2;i++){
            if (sum[initial[i]]>Max){
                Max=sum[initial[i]];
                cnt_i=initial[i];
            }
        }
        return cnt_i;
    }
};