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

class Solution {
public:
    vector<int> fa;     // 每个连通分量的根节点
    vector<int> countPerson;    // 每个联通分量的总人数
    vector<int> countInitial;   // 每个联通分量内init感染数
    vector<int> isInitial;      // 记录initial的下标
    vector<vector<int>> graph;
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
        int n = graph.size();
        fa = vector<int>(n);
        countPerson = vector<int>(n, 1);
        countInitial = vector<int>(n, 0);
        isInitial = vector<int>(n, 0);
        this->graph = graph;

        for(int i = 0; i < n; ++i) {
            fa[i] = i;
        }
        for(int i = 0;i < initial.size(); ++i) {
            isInitial[initial[i]] = 1;
            countInitial[initial[i]] = 1;
        }

        for(int i = 0; i < n; ++i) {
            for(int j = 0;j < n; ++j) {
                if(graph[i][j]) merge(i,j);
            }
        }

        int ret = -1;int maxRemovePerson = 0;
        for(int i = 0;i < initial.size(); ++i) {
            int f = find(initial[i]);
            int removePerson = countPerson[f];
            // 若连通分量内有超过1个的initial节点，则需要取最优
            if(countInitial[f] > 1) {
                removePerson = getRemovePerson(initial[i]) + 1;
            }

            if(maxRemovePerson < removePerson) {
                ret = initial[i];
                maxRemovePerson = removePerson;
            } else if(maxRemovePerson == removePerson && initial[i] < ret) {
                ret = initial[i];
            }
        }
        return ret;
    }

    int find(int x) {
        if(x == fa[x]) return x;
        return fa[x] = find(fa[x]);
    }

    void merge(int x,int y) {
        int fa_x = find(x);
        int fa_y = find(y);
        if(fa_x != fa_y) {
            fa[fa_y] = fa_x;
            countPerson[fa_x] += countPerson[fa_y];
            countInitial[fa_x] += countInitial[fa_y];
        }
    }

    // 获取到去除x结点后避免被感染的结点数量
    int getRemovePerson(int x) {
        int n = graph.size();
        bool isLoop = findLoop(x);  // 查询当前根节点是否在环中
        // 若当前节点在环中，则表示即使去除当前结点，整个连通分量还是会被全部感染
        if(isLoop) return 0;        // 是，则直接退出

        // 若不是，则以当前结点为根节点，查找图的结点数量
        int count = 0;              // 记录结点数
        vector<int> vis(n, 0);      // 记录结点是否已被访问
        vis[x] = 1;
        for(int i = 0;i < n; ++i) {
            if(i != x && graph[x][i] && !vis[i]) count += countChildPerson(i, vis);
        }
        return count;
    }


    int countChildPerson(int x, vector<int>& vis) {
        queue<int> q;
        q.push(x);
        vis[x] = 1;
        int count = 0;
        bool status = false;    // 记录路径中是否有被感染的结点
        while(!q.empty()) {
            auto cur = q.front();
            q.pop();
            if(isInitial[cur]) status = true;   // cur为感染结点
            count++;
            for(int i = 0; i < graph.size(); ++i) {
                if(graph[cur][i] && !vis[i]) {
                    vis[i] = 1;
                    q.push(i);
                }
            }
        }
        if(status) return 0;    // 若路径中存在感染结点，则整条路径都会被感染
        return count;
    }

    // 判断x是否为循环结点
    bool findLoop(int x) {
        int n = graph.size();
        vector<int> vis(n, 0);
        bool isHasOutNode = false; // 判断当前结点是否有超过1个的外接结点
        bool status = false;

        vis[x] = 1;

        for(int i = 0; i < n; ++i) {
            if(i == x) continue;
            if(graph[x][i] && !vis[i]) {
                if(isHasOutNode) return false;  // 若有外接结点，则非环
                isHasOutNode = true;
                vis[i] = 1;
                status = dfs(x, x, i, vis); // dfs查询x是否在环中
            }
        }
        return status;
    }

    bool dfs(int target, int fa, int x, vector<int>& vis) {

        for(int i = 0; i < graph.size(); ++i) {
            if(i == x || i == fa) continue; // 若为当前结点或是父节点则跳过
            if(graph[x][i]) {
                if(vis[i] && i == target) return true; // 若找到target，则表示target在环中
                if(!vis[i]) {
                    vis[i] = 1;
                    bool status = dfs(target, x, i, vis);
                    if(status) return true;
                }
            }
        }
        return false;
    }
};