LeetCode 928

class Solution {
public:
    int find(int x){   //返回x的父节点
        int t = x;
        while(parent[x]>=0){  //x是跟节点
            x = parent[x];
        }
        //压秩
        // int next;
        // while(t!=x){
        //     next = parent[t];
        //     parent[t] = x;
        //     t = next;
        // }
        return x;
    }
    
    void Union(int x,int y){    //合并两个集合
        int px=find(x),py=find(y);
        //少的合并到多的上面
        if(parent[px] < parent[py]){  //x多
            //px作为跟
            parent[px] += parent[py];
            parent[py] = px;
        }
        else{   //y多
            parent[py] += parent[px];
            parent[px] = py;
        }
    }
    
    vector<int> parent;
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
        int n = graph.size(); //一共有多少节点
        parent.assign(n,-1);  //赋初值
        vector<bool> outside(n,1);
        //outside[i]==1, not in initial
        for(auto u:initial)
            outside[u] = false;
        //只有两者都在intial之外时，才进行合并
        for(int i=0;i<n;i++){
            if(outside[i]){
                for(int j=i+1;j<n;j++){
                    if(outside[j]){
                        if(graph[i][j]){
                            Union(i,j);
                        }
                    }
                }
            }
        }
        
        vector<int> cnt(n,0);  //当前节点一共能被多少个节点所影响
        unordered_map<int,set<int>> mp;  //源节点->被影响的节点集合
        for(auto i:initial){    //枚举源节点
            set<int> st;  //能被u所影响的节点集合
            for(int j=0;j<n;j++){
                if(outside[j] && graph[i][j]){
                    st.insert(find(j));
                }
            }
            
            mp[i] = st;
            for(auto u:st){
                cnt[u]++;  //被initial中影响的个数
            }
        }
        
        //由于连同分量之间是互相连同的，只要影响到一个，互相之间都会产生影响，因此只需要统计每个连同分量之间被影响的个数
        int resIdx=-1;
        int resSize = -1;
        auto it = mp.begin();
        for(;it!=mp.end();it++){
            int u = it->first;  //源节点   
            auto st = it->second;   //被影响的节点集合
            int size = 0;
            for(auto v:st){
                if(cnt[v]==1){
                    size += -parent[v];
                }
            }
            
            if(size>resSize || (size==resSize && u<resIdx)){
                resSize = size;
                resIdx = u;
            }
        }
        
        return resIdx;
    }
};