Union-Find (Disjoint Set Union)

// Template 1: union-find (disjoint set union) over items 0 … n - 1.
// parent[] stores a forest: following parent links from any item ends at its group's root.
// Two items are in the same group exactly when they have the same root.
export class UnionFind {
  constructor(n) {
    this.parent = Array.from({ length: n }, (_, i) => i); // every item starts as its own root
    this.rank = new Array(n).fill(0);                     // upper bound on each root's tree height
    this.count = n;                                       // number of separate groups
  }

  find(x) {
    let root = x;
    while (this.parent[root] !== root) root = this.parent[root];
    while (this.parent[x] !== root) {                     // path compression: point the whole path at the root
      const next = this.parent[x];
      this.parent[x] = root;
      x = next;
    }
    return root;
  }

  union(a, b) {
    let ra = this.find(a);
    let rb = this.find(b);
    if (ra === rb) return false;                          // already in one group (an edge here closes a cycle)
    if (this.rank[ra] < this.rank[rb]) [ra, rb] = [rb, ra];
    this.parent[rb] = ra;                                 // union by rank: the shorter tree goes under the taller
    if (this.rank[ra] === this.rank[rb]) this.rank[ra]++;
    this.count--;
    return true;
  }

  connected(a, b) {
    return this.find(a) === this.find(b);
  }
}

import { UnionFind } from './union-find.js';

// Template 2: count groups. Start with n groups (every item alone); each union that really
// merges two roots removes one group. Pairs inside an existing group change nothing.
export function countGroups(n, pairs) {
  const uf = new UnionFind(n);
  for (const [a, b] of pairs) uf.union(a, b);
  return uf.count;
}

// The same loop answers "how big is each group?": bucket every item under its root.
export function groupsOf(n, pairs) {
  const uf = new UnionFind(n);
  for (const [a, b] of pairs) uf.union(a, b);
  const groups = new Map();
  for (let i = 0; i < n; i++) {
    const root = uf.find(i);        // always find(i), never parent[i]: a parent is not always a root
    if (!groups.has(root)) groups.set(root, []);
    groups.get(root).push(i);
  }
  return [...groups.values()];
}

import { UnionFind } from './union-find.js';

// Template 3: cycle detection in an undirected graph. Add the edges one by one; the first edge
// whose two ends already share a root closes a cycle (they were connected without it).
export function firstCycleEdge(n, edges) {
  const uf = new UnionFind(n);
  for (const [a, b] of edges) {
    if (!uf.union(a, b)) return [a, b];
  }
  return null; // the edges form a forest
}

// A graph is a tree when it has exactly n - 1 edges and none of them closes a cycle.
export function isTree(n, edges) {
  return edges.length === n - 1 && firstCycleEdge(n, edges) === null;
}

// Template 4: union-find over arbitrary keys (emails, words, "r,c" cells) with union by size.
// Keys are added the first time they are seen, so nothing has to be numbered up front,
// and every root knows how many items its group holds.
export class KeyedUnionFind {
  constructor() {
    this.parent = new Map();
    this.sizes = new Map(); // root -> number of items in its group
  }

  add(x) {
    if (!this.parent.has(x)) {
      this.parent.set(x, x);
      this.sizes.set(x, 1);
    }
  }

  find(x) {
    this.add(x);
    let root = x;
    while (this.parent.get(root) !== root) root = this.parent.get(root);
    while (this.parent.get(x) !== root) {   // path compression
      const next = this.parent.get(x);
      this.parent.set(x, root);
      x = next;
    }
    return root;
  }

  union(a, b) {
    let ra = this.find(a);
    let rb = this.find(b);
    if (ra === rb) return false;
    if (this.sizes.get(ra) < this.sizes.get(rb)) [ra, rb] = [rb, ra];
    this.parent.set(rb, ra);                // union by size: the smaller group joins the larger
    this.sizes.set(ra, this.sizes.get(ra) + this.sizes.get(rb));
    return true;
  }

  sizeOf(x) {
    return this.sizes.get(this.find(x));
  }

  groups() {
    const out = new Map();
    for (const x of this.parent.keys()) {
      const root = this.find(x);
      if (!out.has(root)) out.set(root, []);
      out.get(root).push(x);
    }
    return [...out.values()];
  }
}