WeightedQuickUnionFind.js

class WeightedQuickUnionFind {
  constructor (ids) {
    this._map = new Map() // id-to-component map
    this._sizes = new Map() // subtree size map
    this.length = ids.length // component count

    ids.forEach((id) => {
      this._map.set(id, id)
      this._sizes.set(id, 1)
    })
  }

  // are `idA` and `idB` in the same component
  connected = (idA, idB) =>
    this._root(idA) === this._root(idB)

  // size of the component to which `id` belongs
  size = (id) =>
    this._sizes.get(this._root(id))

  // merges the components containing `idA` and `idB`
  union (idA, idB) {
    const rootA = this._root(idA)
    const rootB = this._root(idB)

    if (rootA === rootB) return // exit early

    const sizeA = this._sizes.get(rootA)
    const sizeB = this._sizes.get(rootB)

    if (sizeA < sizeB) {
      this._map.set(rootA, rootB)
      this._sizes.set(rootB, sizeB + sizeA)
    } else {
      this._map.set(rootB, rootA)
      this._sizes.set(rootA, sizeA + sizeB)
    }

    this.length--
  }

  _root (id) {
    let root = id
    while (root !== this._map.get(root)) {
      root = this._map.get(root)
    }
    // path compression
    while (id !== root) {
      const nextId = this._map.get(id)
      this._map.set(id, root)
      id = nextId
    }
    return root
  }
}