Minimize Malware Spread

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

In a network of nodes, each node i is directly connected to another node j if and only if graph[i][j] = 1.

Some nodes initial are initially infected by malware. Whenever two nodes are directly connected and at least one of those two nodes is infected by malware, both nodes will be infected by malware. This spread of malware will continue until no more nodes can be infected in this manner.

Suppose M(initial) is the final number of nodes infected with malware in the entire network, after the spread of malware stops.

We will remove one node from the initial list. Return the node that if removed, would minimize M(initial). If multiple nodes could be removed to minimize M(initial), return such a node with the smallest index.

Note that if a node was removed from the initial list of infected nodes, it may still be infected later as a result of the malware spread.

Example 1:

Input: graph = [[1,1,0],[1,1,0],[0,0,1]], initial = [0,1]

Output: 0

Example 2:

Input: graph = [[1,0,0],[0,1,0],[0,0,1]], initial = [0,2]

Output: 0

Example 3:

Input: graph = [[1,1,1],[1,1,1],[1,1,1]], initial = [1,2]

Output: 1

Note:

1 < graph.length = graph[0].length <= 300

0 <= graph[i][j] == graph[j][i] <= 1

graph[i][i] = 1

1 <= initial.length < graph.length

0 <= initial[i] < graph.length

Thoughts

当有两个感染节点处于同一个component, 那无论删除哪个节点最后感染节点数都只会少1。只有当一个component中只有单独的感染时，移除它会使得整个component都不被感染。 图中判断哪些元素时同一个component想到了uf, 并且记下每个component的size.

Code

class Solution {
public:
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
        const auto n = graph.size();
        for (int i = 0; i < n; ++i) parents.push_back(i);
        for (int i = 0; i < n; ++i) {
            for (int j = i + 1; j < n; ++j) {
                if (graph[i][j]) un(i, j);
            }
        }

        // The size of each component; how many malwares in each component
        vector<int> size(n, 0), malware(n, 0);
        for (int i = 0; i < n; ++i) ++size[find(i)]; 
        for (int i : initial) ++malware[find(i)];
        pair<int, int> res = {1, 0};
        for (int i : initial) {
            res = min(res, {(malware[find(i)] == 1) * (-size[find(i)]), i});
        }
        return res.second;
    }

private:
    vector<int> parents;
    int find(int x) {
        if (x != parents[x]) {
            parents[x] = find(parents[x]);
        }
        return parents[x];
    }

    void un(int x, int y) {
        x = find(x), y = find(y);
        if (x != y) {
            parents[x] = y;
        }
    }
};

Analysis

TC: O(N^2)

每次union操作时间复杂度时amortized O(1).