LeetCode Solutions
924. Minimize Malware Spread
Time: $O(n^2)$, where $n = |\texttt{graph}|$ Space: $O(n)$
class UnionFind {
public:
UnionFind(int n) : id(n) {
iota(begin(id), end(id), 0);
}
void union_(int u, int v) {
id[find(u)] = find(v);
}
int find(int u) {
return id[u] == u ? u : id[u] = find(id[u]);
}
private:
vector<int> id;
};
class Solution {
public:
int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
const int n = graph.size();
UnionFind uf(n);
vector<int> ufSize(n);
vector<int> malwareCount(n);
for (int i = 0; i < n; ++i)
for (int j = i + 1; j < n; ++j)
if (graph[i][j] == 1)
uf.union_(i, j);
for (int i = 0; i < n; ++i)
++ufSize[uf.find(i)];
for (const int i : initial)
++malwareCount[uf.find(i)];
sort(begin(initial), end(initial));
int ans = initial[0];
int maxUfSize = 0;
// Find the max 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;
}
};
class UnionFind {
public UnionFind(int n) {
id = new int[n];
for (int i = 0; i < n; ++i)
id[i] = i;
}
public void union(int u, int v) {
id[find(u)] = find(v);
}
public int find(int u) {
return id[u] == u ? u : (id[u] = find(id[u]));
}
private int[] id;
}
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.union(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 max 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;
}
}