Minimize Malware Spread LeetCode solution in JavaScript

1
var minMalwareSpread = function (graph, initial) {
2
  let n = graph.length,
3
    uf = new UnionFind(n);
4
  for (let i = 0; i < n; i++) {
5
    for (let j = i + 1; j < n; j++) {
6
      if (graph[i][j]) uf.union(i, j);
7
    }
8
  }
9

10
  // Get number of nodes in each connected component
11
  let nodesCount = {};
12
  for (let i = 0; i < n; i++) {
13
    let parent = uf.find(i);
14
    nodesCount[parent] = (nodesCount[parent] || 0) + 1;
15
  }
16

17
  // Find the initial node connected to the most amount of other nodes, where it is the only malware infected node in its group.
18
  let max = 0,
19
    res = Infinity;
20
  for (let node of initial) {
21
    let malwareNodes = malwareConnected(node);
22
    let parent = uf.find(node);
23
    let nodesAffected = malwareNodes === 1 ? nodesCount[parent] : 0;
24
    if (nodesAffected > max || (nodesAffected === max && node < res)) {
25
      res = node;
26
      max = nodesAffected;
27
    }
28
  }
29
  return res;
30

31
  function malwareConnected(node) {
32
    // gets number of malware infected nodes connected to node
33
    let count = 0;
34
    for (let malware of initial) {
35
      if (uf.isConnected(node, malware)) count++;
36
    }
37
    return count;
38
  }
39
};
40

41
class UnionFind {
42
  constructor(size) {
43
    this.rank = Array(size);
44
    this.root = Array(size);
45
    for (let i = 0; i < size; i++) {
46
      this.rank[i] = 1;
47
      this.root[i] = i;
48
    }
49
  }
50
  find(x) {
51
    if (this.root[x] === x) return x;
52
    return (this.root[x] = this.find(this.root[x]));
53
  }
54
  union(x, y) {
55
    let rootX = this.find(x),
56
      rootY = this.find(y);
57
    if (rootX === rootY) return false;
58
    if (this.rank[rootX] < this.rank[rootY]) this.root[rootX] = rootY;
59
    else if (this.rank[rootX] > this.rank[rootY]) this.root[rootY] = rootX;
60
    else {
61
      this.root[rootY] = rootX;
62
      this.rank[rootX]++;
63
    }
64
    return true;
65
  }
66
  isConnected(x, y) {
67
    return this.find(x) === this.find(y);
68
  }
69
}