928. 盡量減少惡意軟體的傳播 II

思路

給定一個鄰接矩陣，代表節點之間連接的情況（無向圖）。
然後給定一個initial陣列，代表有病毒的節點，這些節點會將病毒傳遞出去。
問刪除哪一個病毒節點，能拯救最多節點？

如果一個節點能夠被拯救，那它要恰好被一個病毒所感染，一旦被兩個或以上的病毒感染，不管移除哪一個節點，該節點都無法拯救，反之，完全沒有被感染，就不需要被拯救，也不會拯救數量中。

1. 先將沒有病毒的節點，合併成一些集合。
2. 再遍歷病毒，把它會感染到的集合標記起來，記錄在infect陣列中。

• 完全沒有被感染，保持初始值-1。
• 一個集合被恰好感染一次，將感染的病毒作為標籤紀錄。
• 一個集合被感染了多次以上，那這整個集合就沒救了，標為-2。

3. 遍歷根結點，紀錄病毒感染到的陣列大小總和到save陣列。
4. 遍歷save陣列，找出最優解。

程式碼

並查集

class Solution {
private:
    static const int MX = 301;
    int parent[MX];
    int setSize[MX];
    bool virus[MX]; // 記錄點是否被感染
    int save[MX]; // 將根結點刪掉之後，可以拯救多少節點
    int infect[MX]; // infect[i] = -1, 目前這個集合沒有受到感染
    // infect[i] >= 0, 目前這個集合被 infect[i] 感染。
    // infect[i] = -2, 被多個感染源感染，沒救了。
    void build(int n, vector<int>& initial) {
        iota(parent, parent + n, 0);
        fill(setSize, setSize + n, 1);
        fill(virus, virus + n, false);
        fill(save, save + n, 0);
        fill(infect, infect + n, -1);
        for(int x : initial) {
            virus[x] = true;
        }
    }
    int find(int x) {
        if(x != parent[x]) {
            parent[x] = find(parent[x]);
        }
        return parent[x];
    }
    void unite(int x, int y) {
        int rootX = find(x);
        int rootY = find(y);
        if(rootX != rootY) {
            parent[rootX] = rootY;
            setSize[rootY] += setSize[rootX];
        }
    }
public:
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
        int n = graph.size();
        build(n, initial);
        for(int i = 0; i < n; i++) {
            for(int j = 0; j < n; j++) {
                // 把沒有病毒的部分合併起來
                if(!virus[i] && !virus[j] && graph[i][j] == 1) { 
                    unite(i, j);
                }
            }
        }
        for(int x : initial) { // 病毒的鄰居，會被病毒所感染
            for(int j = 0; j < n; j++) {
                if(x != j && !virus[j] && graph[x][j] == 1) {
                    int& temp = infect[find(j)]; // 這裡是引用
                    if(temp == -1) {
                        temp = x;
                    }
                    else if(temp != -2 && temp != x) { // 已受感染，且感染的病毒不相同
                        temp = -2; // 宣告不治
                    }
                }
            }
        }
        for(int i = 0; i < n; i++) {
            // i == find(i) 代表它是集合的根結點，只有根結點才能代表集合。
            if(i == find(i) && infect[i] >= 0) {
                save[infect[i]] += setSize[i];
            }
        }
        ranges::sort(initial); // 也可以不排序，不過下面遍歷就要多加判斷
        int res = initial[0];
        int mxSave = save[initial[0]];
        for(int x : initial) {
            if(save[x] > mxSave) {
                res = x;
                mxSave = save[x];
            }
        }
        return res;
    }
};

複雜度分析

• 時間複雜度：$O(n^2)$
• 空間複雜度：$O(n)$