Minimize Malware Spread - Union Find [JS]

Description 

Solution: Union Find

It's only worth removing nodes that are not connected to other malware infected nodes.

1. Use union find to connect each node based on the graph.
2. Get the number of nodes in each connected component (group by the parent nodes).
3. Find the initial node connected to the most amount of other nodes, where it is the only malware infected node in its group.
   Count the number of malware infected nodes it is connected to.
   If this count is > 1, it is useless to remove this node. Otherwise, we save the whole group.

Time Complexity: O(n^2) 174ms
Space Complexity: O(n) 52.6MB

var minMalwareSpread = function(graph, initial) {
  let n = graph.length, uf = new UnionFind(n);
  for (let i = 0; i < n; i++) {
    for (let j = i + 1; j < n; j++) {
      if (graph[i][j]) uf.union(i, j);
    }
  }
  
  // Get number of nodes in each connected component
  let nodesCount = {};
  for (let i = 0; i < n; i++) {
    let parent = uf.find(i);
    nodesCount[parent] = (nodesCount[parent] || 0) + 1;
  }
  
  // Find the initial node connected to the most amount of other nodes, where it is the only malware infected node in its group.
  let max = 0, res = Infinity;
  for (let node of initial) {
    let malwareNodes = malwareConnected(node);
    let parent = uf.find(node);
    let nodesAffected = malwareNodes === 1 ? nodesCount[parent] : 0;
    if (nodesAffected > max || (nodesAffected === max && node < res)) {
      res = node;
      max = nodesAffected;
    }
  }
  return res;
  
  function malwareConnected(node) { // gets number of malware infected nodes connected to node
    let count = 0;
    for (let malware of initial) {
      if (uf.isConnected(node, malware)) count++;
    }
    return count;
  }
};

class UnionFind {
  constructor(size) {
    this.rank = Array(size);
    this.root = Array(size);
    for (let i = 0; i < size; i++) {
      this.rank[i] = 1;
      this.root[i] = i;
    }
  }
  find(x) {
    if (this.root[x] === x) return x;
    return this.root[x] = this.find(this.root[x]);
  }
  union(x, y) {
    let rootX = this.find(x), rootY = this.find(y);
    if (rootX === rootY) return false;
    if (this.rank[rootX] < this.rank[rootY]) this.root[rootX] = rootY;
    else if (this.rank[rootX] > this.rank[rootY]) this.root[rootY] = rootX;
    else {
      this.root[rootY] = rootX;
      this.rank[rootX]++;
    }
    return true;
  }
  isConnected(x, y) {
    return this.find(x) === this.find(y);
  }
}

javascript