Solution.cpp

class Solution {
public:
    vector<int> p;

    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
        int n = graph.size();
        p.resize(n);
        for (int i = 0; i < n; ++i) p[i] = i;
        vector<int> size(n, 1);
        for (int i = 0; i < n; ++i) {
            for (int j = i + 1; j < n; ++j) {
                if (graph[i][j]) {
                    int pa = find(i), pb = find(j);
                    if (pa == pb) continue;
                    p[pa] = pb;
                    size[pb] += size[pa];
                }
            }
        }
        int mi = 400;
        int res = initial[0];
        sort(initial.begin(), initial.end());
        for (int i = 0; i < initial.size(); ++i) {
            int t = 0;
            unordered_set<int> s;
            for (int j = 0; j < initial.size(); ++j) {
                if (i == j) continue;
                if (s.count(find(initial[j]))) continue;
                s.insert(find(initial[j]));
                t += size[find(initial[j])];
            }
            if (mi > t) {
                mi = t;
                res = initial[i];
            }
        }
        return res;
    }

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