Minimize Malware Spread.js

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);
  }
}