• Time: $O(n^2)$, where $n = |\texttt{graph}|$
• Space: $O(n)$
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 class UnionFind { public: UnionFind(int n) : id(n), rank(n) { iota(id.begin(), id.end(), 0); } void unionByRank(int u, int v) { const int i = find(u); const int j = find(v); if (i == j) return; if (rank[i] < rank[j]) { id[i] = j; } else if (rank[i] > rank[j]) { id[j] = i; } else { id[i] = j; ++rank[j]; } } int find(int u) { return id[u] == u ? u : id[u] = find(id[u]); } private: vector id; vector rank; }; class Solution { public: int minMalwareSpread(vector>& graph, vector& initial) { const int n = graph.size(); UnionFind uf(n); vector ufSize(n); vector malwareCount(n); for (int i = 0; i < n; ++i) for (int j = i + 1; j < n; ++j) if (graph[i][j] == 1) uf.unionByRank(i, j); for (int i = 0; i < n; ++i) ++ufSize[uf.find(i)]; for (const int i : initial) ++malwareCount[uf.find(i)]; ranges::sort(initial); int ans = initial[0]; int maxUfSize = 0; // Find the maximum union's malware if it only contains 1 malware. for (const int i : initial) { const int id = uf.find(i); if (ufSize[id] > maxUfSize && malwareCount[id] == 1) { maxUfSize = ufSize[id]; ans = i; } } return ans; } }; 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 class UnionFind { public UnionFind(int n) { id = new int[n]; rank = new int[n]; for (int i = 0; i < n; ++i) id[i] = i; } public void unionByRank(int u, int v) { final int i = find(u); final int j = find(v); if (i == j) return; if (rank[i] < rank[j]) { id[i] = j; } else if (rank[i] > rank[j]) { id[j] = i; } else { id[i] = j; ++rank[j]; } } public int find(int u) { return id[u] == u ? u : (id[u] = find(id[u])); } private int[] id; private int[] rank; } class Solution { public int minMalwareSpread(int[][] graph, int[] initial) { final int n = graph.length; UnionFind uf = new UnionFind(n); int[] ufSize = new int[n]; int[] malwareCount = new int[n]; for (int i = 0; i < n; ++i) for (int j = i + 1; j < n; ++j) if (graph[i][j] == 1) uf.unionByRank(i, j); for (int i = 0; i < n; ++i) ++ufSize[uf.find(i)]; for (final int i : initial) ++malwareCount[uf.find(i)]; Arrays.sort(initial); int ans = initial[0]; int maxUfSize = 0; // Find the maximum union's malware if it only contains 1 malware. for (final int i : initial) { final int id = uf.find(i); if (ufSize[id] > maxUfSize && malwareCount[id] == 1) { maxUfSize = ufSize[id]; ans = i; } } return ans; } }