The Union-Find With Rank I Copy Into Every Contest

// DSU (Disjoint Set Union) / Union-Find. The two optimizations that take the
// per-operation cost from O(n) to amortized O(alpha(n)) (effectively O(1)):
//   1. Path compression in find(): every traversal flattens the chain.
//   2. Union-by-rank: attach the shorter tree under the taller one.
// Iterative find() avoids stack overflow on chains of 1e6 nodes.

class DSU {
    constructor(n) {
        this.parent = new Array(n);
        this.rank = new Array(n).fill(0);
        this.componentCount = n;
        for (let i = 0; i < n; i++) this.parent[i] = i;
    }

    find(x) {
        // Two-pass: first walk to the root, then point every node along the path at the root.
        let root = x;
        while (this.parent[root] !== root) root = this.parent[root];
        let cur = x;
        while (this.parent[cur] !== root) {
            const next = this.parent[cur];
            this.parent[cur] = root;
            cur = next;
        }
        return root;
    }

    union(a, b) {
        const ra = this.find(a);
        const rb = this.find(b);
        if (ra === rb) return false;  // already connected
        if (this.rank[ra] < this.rank[rb]) {
            this.parent[ra] = rb;
        } else if (this.rank[ra] > this.rank[rb]) {
            this.parent[rb] = ra;
        } else {
            this.parent[rb] = ra;
            this.rank[ra] += 1;
        }
        this.componentCount -= 1;
        return true;
    }

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

const dsu = new DSU(6);
console.log('start:', dsu.componentCount);  // 6
dsu.union(0, 1);
dsu.union(2, 3);
dsu.union(4, 5);
console.log('after 3 unions:', dsu.componentCount);  // 3
dsu.union(1, 3);
dsu.union(3, 5);
console.log('all connected:', dsu.componentCount);   // 1
console.log('connected(0, 5)?', dsu.connected(0, 5));  // true
console.log('union(0, 5) returns:', dsu.union(0, 5));  // false (already connected)

// The canonical use of DSU: sort edges by weight, add each in order if its
// endpoints are not yet connected. The DSU answers 'are these in the same
// component?' in amortized O(1).

class DSU {
    constructor(n) {
        this.parent = Array.from({ length: n }, (_, i) => i);
        this.rank = new Array(n).fill(0);
        this.componentCount = n;
    }
    find(x) {
        let r = x; while (this.parent[r] !== r) r = this.parent[r];
        while (this.parent[x] !== r) { const n = this.parent[x]; this.parent[x] = r; x = n; }
        return r;
    }
    union(a, b) {
        const ra = this.find(a), rb = this.find(b);
        if (ra === rb) return false;
        if (this.rank[ra] < this.rank[rb]) this.parent[ra] = rb;
        else if (this.rank[ra] > this.rank[rb]) this.parent[rb] = ra;
        else { this.parent[rb] = ra; this.rank[ra] += 1; }
        this.componentCount -= 1;
        return true;
    }
}

function kruskal(n, edges) {
    const sorted = [...edges].sort((a, b) => a[2] - b[2]);
    const dsu = new DSU(n);
    let cost = 0;
    const tree = [];
    for (const [u, v, w] of sorted) {
        if (dsu.union(u, v)) {
            cost += w;
            tree.push([u, v, w]);
            if (tree.length === n - 1) break;
        }
    }
    return { cost, tree, connected: dsu.componentCount === 1 };
}

const edges = [
    [0, 1, 4], [0, 2, 3], [1, 2, 1], [1, 3, 2], [2, 3, 4], [3, 4, 2], [4, 5, 6], [3, 5, 7],
];
const result = kruskal(6, edges);
console.log('MST cost:', result.cost);
console.log('MST edges:', result.tree);
console.log('graph fully connected?', result.connected);

// Two variations I keep separately:
//   - Union-by-size: same idea as union-by-rank but tracks subtree size.
//     Useful when problems ask 'how big is the component containing x?'.
//   - DSU with weights: each node carries its 'distance' from the root in
//     some semiring (offset, parity, ratio). Used for problems like
//     'do these expressions agree on a relative value?'.

class SizedDSU {
    constructor(n) {
        this.parent = Array.from({ length: n }, (_, i) => i);
        this.size = new Array(n).fill(1);
    }
    find(x) {
        let r = x; while (this.parent[r] !== r) r = this.parent[r];
        while (this.parent[x] !== r) { const n = this.parent[x]; this.parent[x] = r; x = n; }
        return r;
    }
    union(a, b) {
        const ra = this.find(a), rb = this.find(b);
        if (ra === rb) return false;
        if (this.size[ra] < this.size[rb]) {
            this.parent[ra] = rb; this.size[rb] += this.size[ra];
        } else {
            this.parent[rb] = ra; this.size[ra] += this.size[rb];
        }
        return true;
    }
    componentSize(x) { return this.size[this.find(x)]; }
}

const sd = new SizedDSU(8);
for (const [a, b] of [[0, 1], [1, 2], [3, 4], [5, 6], [6, 7]]) sd.union(a, b);
console.log('size of component(0):', sd.componentSize(0));  // 3
console.log('size of component(3):', sd.componentSize(3));  // 2
console.log('size of component(5):', sd.componentSize(5));  // 3

// Sketch of weighted DSU; the production version threads a numeric 'weight'
// alongside parent and updates it during find/union so 'distance to root'
// stays consistent. We omit the full implementation here because the right
// weight depends on the problem (XOR, addition modulo k, ratio).
console.log('weighted DSU: tracks distance-to-root; pick the weight per problem.');