Simple Disjoint Set

DisjointSet.js

function DJSet() {
    this.nodes = new Map(); // id -> Node
    this.sets = new Map();  // id -> Node
}

function Node(id) {
    this.id = id;
    this.rank = 0;
    this.parent = this;
}

DJSet.prototype.getOrCreateSet = function(id) {
    if (this.nodes.has(id)) {
      return this.nodes.get(id);
    }
    
    const node = new Node(id);
    this.nodes.set(id, node);
    this.sets.set(id, node);

    return node;
}

DJSet.prototype.findSet = function(id) {
    if (!this.nodes.has(id)) {
        return null;
    }
    
    let node = this.nodes.get(id);
    let current = node;
    
    while (current !== current.parent) {
        current = current.parent;
    }
    
    node.parent = current;
    
    return current;
}

DJSet.prototype.union = function(id1, id2) {
    const node1 = this.findSet(id1);
    const node2 = this.findSet(id2);

    if (!node1) {
      throw Error(`Set "${id1}" not exist.`)
    }

    if (!node2) {
      throw Error(`Set "${id2}" not exist.`)
    }

    if (node1.rank > node2.rank) {
      node2.parent = node1;
      node1.rank += 1;
      this.sets.delete(node2.id);

      return node1;
    } else {
      node1.parent = node2;
      node2.rank += 1;
      this.sets.delete(node1.id)

      return node2;
    }
}


const tree = new DJSet();

const node1 = tree.getOrCreateSet(1);
const node2 = tree.getOrCreateSet(2);
const node3 = tree.getOrCreateSet(3);
const node4 = tree.getOrCreateSet(4);
const node5 = tree.getOrCreateSet(5);

tree.union(1, 2);
tree.union(3, 4);
tree.union(1, 4);

console.assert(tree.nodes.size === 5);
console.assert(tree.sets.size === 2);

console.assert(tree.findSet(1) === tree.findSet(2));
console.assert(tree.findSet(2) === tree.findSet(3));
console.assert(tree.findSet(3) === tree.findSet(4));
console.assert(tree.findSet(5) === node5);

https://repl.it/@SHND1/Disjoint-Set