Union-Find Disjoint Set Explained — Path Compression, Union by Rank & Real-World Use Cases

/**
 * Naive Union-Find — NO optimisations.
 * This is the baseline. Study the pathological worst case.
 * DO NOT use this in production.
 */
public class NaiveUnionFind {

    private final int[] parent; // parent[i] = direct parent of node i
    private int componentCount;  // how many separate components exist right now

    public NaiveUnionFind(int totalNodes) {
        this.parent = new int[totalNodes];
        this.componentCount = totalNodes;

        // Initially every node is its own root — n separate components
        for (int node = 0; node < totalNodes; node++) {
            parent[node] = node;
        }
    }

    /**
     * Walk up the parent chain until we reach the root.
     * Root is identified by: parent[root] == root.
     * WORST CASE: O(n) if the tree has degenerated into a chain.
     */
    public int find(int node) {
        while (parent[node] != node) {
            node = parent[node]; // climb one level toward the root
        }
        return node;
    }

    /**
     * Merge the components containing nodeA and nodeB.
     * Returns false if they were already in the same component.
     */
    public boolean union(int nodeA, int nodeB) {
        int rootA = find(nodeA);
        int rootB = find(nodeB);

        if (rootA == rootB) {
            return false; // already connected — unioning would create a cycle
        }

        // Naive: always make rootB point to rootA — no rank consideration
        parent[rootB] = rootA;
        componentCount--;
        return true;
    }

    public boolean connected(int nodeA, int nodeB) {
        return find(nodeA) == find(nodeB);
    }

    public int getComponentCount() {
        return componentCount;
    }

    public static void main(String[] args) {
        // Simulate: 6 cities, progressively connected
        NaiveUnionFind cities = new NaiveUnionFind(6);

        System.out.println("Initial components: " + cities.getComponentCount()); // 6

        cities.union(0, 1); // city 0 and 1 share a road
        cities.union(1, 2); // city 1 and 2 share a road — degenerate chain starts here
        cities.union(2, 3);
        cities.union(3, 4);

        System.out.println("After 4 unions, components: " + cities.getComponentCount()); // 2
        System.out.println("City 0 and 4 connected? " + cities.connected(0, 4));  // true
        System.out.println("City 0 and 5 connected? " + cities.connected(0, 5));  // false

        // The tree for nodes 0-4 is now a chain: 4 -> 3 -> 2 -> 1 -> 0
        // find(4) must climb 4 levels — O(n) already with just 5 nodes
        System.out.println("Root of city 4: " + cities.find(4)); // 0
    }
}

/**
 * Production-grade Union-Find with:
 *   1. Path Compression  — flattens trees during find()
 *   2. Union by Rank     — keeps trees shallow during union()
 *
 * Amortised time per operation: O(α(n)) — effectively O(1)
 * Space: O(n)
 */
public class OptimisedUnionFind {

    private final int[] parent;
    private final int[] rank;    // upper bound on subtree height
    private final int[] size;    // number of nodes in the component (bonus for free)
    private int componentCount;

    public OptimisedUnionFind(int totalNodes) {
        parent         = new int[totalNodes];
        rank           = new int[totalNodes];
        size           = new int[totalNodes];
        componentCount = totalNodes;

        for (int node = 0; node < totalNodes; node++) {
            parent[node] = node; // every node starts as its own root
            rank[node]   = 0;    // single-node tree has height 0
            size[node]   = 1;    // each component starts with 1 node
        }
    }

    /**
     * PATH COMPRESSION — recursive variant (two-pass compression).
     *
     * On the way DOWN the recursion we find the root.
     * On the way BACK UP we rewire every node to point directly to the root.
     * After one call, the entire path is depth-1.
     */
    public int find(int node) {
        if (parent[node] != node) {
            // Recurse to root, then set this node's parent to the root directly
            parent[node] = find(parent[node]); // ← the compression happens here
        }
        return parent[node];
    }

    /**
     * UNION BY RANK.
     *
     * Attach the root with lower rank under the root with higher rank.
     * If ranks are equal, attach either way and increment the surviving root's rank.
     *
     * Returns true if a merge happened, false if they were already connected.
     */
    public boolean union(int nodeA, int nodeB) {
        int rootA = find(nodeA);
        int rootB = find(nodeB);

        if (rootA == rootB) {
            return false; // same component — no merge needed, would form a cycle
        }

        // Attach smaller-rank tree under larger-rank tree
        if (rank[rootA] < rank[rootB]) {
            parent[rootA] = rootB;       // rootA hangs under rootB
            size[rootB]  += size[rootA]; // update component size at the new root
        } else if (rank[rootA] > rank[rootB]) {
            parent[rootB] = rootA;
            size[rootA]  += size[rootB];
        } else {
            // Equal rank: arbitrary choice — attach rootB under rootA
            parent[rootB] = rootA;
            size[rootA]  += size[rootB];
            rank[rootA]++;               // only case where rank increments
        }

        componentCount--;
        return true;
    }

    public boolean connected(int nodeA, int nodeB) {
        return find(nodeA) == find(nodeB);
    }

    /** How many nodes are in the component containing this node? */
    public int componentSize(int node) {
        return size[find(node)];
    }

    public int getComponentCount() {
        return componentCount;
    }

    // ─────────────────────────────────────────────
    // Demonstrate: path compression in action
    // ─────────────────────────────────────────────
    public static void main(String[] args) {
        OptimisedUnionFind dsu = new OptimisedUnionFind(8);

        // Build a chain manually to simulate a worst-case naive tree
        // then show how find() flattens it
        dsu.union(0, 1);
        dsu.union(1, 2);
        dsu.union(2, 3);
        dsu.union(3, 4);
        dsu.union(5, 6);
        dsu.union(6, 7);
        dsu.union(4, 7); // merges two components

        System.out.println("Components remaining: " + dsu.getComponentCount()); // 1
        System.out.println("All 8 nodes connected? " + dsu.connected(0, 7));    // true
        System.out.println("Size of component containing node 3: "
                           + dsu.componentSize(3));                             // 8

        // Verify isolation before connection
        OptimisedUnionFind fresh = new OptimisedUnionFind(5);
        fresh.union(0, 1);
        fresh.union(2, 3);
        System.out.println("Node 1 and 3 connected? " + fresh.connected(1, 3)); // false
        fresh.union(1, 2);
        System.out.println("Node 1 and 3 connected? " + fresh.connected(1, 3)); // true
        System.out.println("Components: " + fresh.getComponentCount());          // 2
    }
}

import java.util.Arrays;

/**
 * Kruskal's Minimum Spanning Tree using Union-Find.
 *
 * Given a weighted undirected graph, find the subset of edges
 * that connects all nodes with the minimum total weight and no cycles.
 *
 * Time:  O(E log E) — dominated by sorting edges
 * Space: O(V)       — for the DSU arrays
 */
public class KruskalMST {

    // ── Inner DSU (self-contained for clarity) ───────────────────────────
    static class DSU {
        private final int[] parent;
        private final int[] rank;

        DSU(int nodeCount) {
            parent = new int[nodeCount];
            rank   = new int[nodeCount];
            for (int i = 0; i < nodeCount; i++) parent[i] = i;
        }

        int find(int node) {
            if (parent[node] != node) parent[node] = find(parent[node]);
            return parent[node];
        }

        boolean union(int a, int b) {
            int ra = find(a), rb = find(b);
            if (ra == rb) return false;             // already connected
            if (rank[ra] < rank[rb]) parent[ra] = rb;
            else if (rank[ra] > rank[rb]) parent[rb] = ra;
            else { parent[rb] = ra; rank[ra]++; }
            return true;
        }
    }

    // ── Edge representation ───────────────────────────────────────────────
    record Edge(int source, int destination, int weight) {}

    // ── Kruskal's algorithm ───────────────────────────────────────────────
    public static int[][] buildMST(int nodeCount, int[][] rawEdges) {
        // Convert raw edge array to Edge objects for clarity
        Edge[] edges = new Edge[rawEdges.length];
        for (int i = 0; i < rawEdges.length; i++) {
            edges[i] = new Edge(rawEdges[i][0], rawEdges[i][1], rawEdges[i][2]);
        }

        // Step 1: Sort edges by weight — greedy choice, cheapest first
        Arrays.sort(edges, (e1, e2) -> e1.weight() - e2.weight());

        DSU dsu            = new DSU(nodeCount);
        int[][] mstEdges   = new int[nodeCount - 1][3]; // MST has exactly n-1 edges
        int edgesAdded     = 0;
        int totalMSTWeight = 0;

        for (Edge edge : edges) {
            if (edgesAdded == nodeCount - 1) break; // MST complete

            // Step 2: Only add edge if it connects two DIFFERENT components
            if (dsu.union(edge.source(), edge.destination())) {
                mstEdges[edgesAdded][0] = edge.source();
                mstEdges[edgesAdded][1] = edge.destination();
                mstEdges[edgesAdded][2] = edge.weight();
                totalMSTWeight += edge.weight();
                edgesAdded++;
                System.out.printf("  Added edge (%d — %d) weight %d%n",
                                  edge.source(), edge.destination(), edge.weight());
            } else {
                System.out.printf("  Skipped edge (%d — %d) weight %d — would form cycle%n",
                                  edge.source(), edge.destination(), edge.weight());
            }
        }

        System.out.println("Total MST weight: " + totalMSTWeight);
        return mstEdges;
    }

    public static void main(String[] args) {
        int nodeCount = 5; // nodes 0, 1, 2, 3, 4

        // Format: {source, destination, weight}
        int[][] edges = {
            {0, 1, 2},
            {0, 3, 6},
            {1, 2, 3},
            {1, 3, 8},
            {1, 4, 5},
            {2, 4, 7},
            {3, 4, 9}
        };

        System.out.println("Building MST for 5-node graph:");
        int[][] mst = buildMST(nodeCount, edges);

        System.out.println("\nMST edges:");
        for (int[] edge : mst) {
            System.out.printf("  %d — %d (weight %d)%n", edge[0], edge[1], edge[2]);
        }
    }
}

/**
 * Iterative (non-recursive) Union-Find with path compression.
 *
 * Solves the StackOverflowError risk of the recursive version.
 * Uses a two-pass iterative approach:
 *   Pass 1 — walk up to find the root.
 *   Pass 2 — walk again and point every node directly to the root.
 *
 * Behaviour is identical to the recursive version but stack-safe.
 */
public class IterativeUnionFind {

    private final int[] parent;
    private final int[] rank;
    private int componentCount;

    public IterativeUnionFind(int totalNodes) {
        parent         = new int[totalNodes];
        rank           = new int[totalNodes];
        componentCount = totalNodes;
        for (int i = 0; i < totalNodes; i++) parent[i] = i;
    }

    /**
     * Iterative path compression — two-pass, stack-safe.
     */
    public int find(int node) {
        // ── Pass 1: find the root ─────────────────────────────────────────
        int root = node;
        while (parent[root] != root) {
            root = parent[root]; // keep climbing
        }

        // ── Pass 2: path compression — rewire everyone to the root ────────
        while (parent[node] != root) {
            int nextNode   = parent[node]; // save next before we overwrite
            parent[node]   = root;         // point directly to root
            node           = nextNode;     // advance up the chain
        }

        return root;
    }

    public boolean union(int nodeA, int nodeB) {
        int rootA = find(nodeA);
        int rootB = find(nodeB);

        if (rootA == rootB) return false;

        if (rank[rootA] < rank[rootB])       { parent[rootA] = rootB; }
        else if (rank[rootA] > rank[rootB])  { parent[rootB] = rootA; }
        else                                 { parent[rootB] = rootA; rank[rootA]++; }

        componentCount--;
        return true;
    }

    public boolean connected(int a, int b) { return find(a) == find(b); }
    public int getComponentCount()         { return componentCount; }

    public static void main(String[] args) {
        // Simulate 1-indexed input (common in interview problems)
        // Nodes labelled 1 to 6 — allocate size 7 (index 0 wasted)
        int maxNode = 6;
        IterativeUnionFind dsu = new IterativeUnionFind(maxNode + 1);

        int[][] connections = {{1,2},{2,3},{4,5},{5,6}};
        for (int[] conn : connections) {
            dsu.union(conn[0], conn[1]);
        }

        System.out.println("Components: " + dsu.getComponentCount()); // 3 (plus unused index 0)
        System.out.println("1 and 3 connected? " + dsu.connected(1, 3)); // true
        System.out.println("1 and 4 connected? " + dsu.connected(1, 4)); // false
        System.out.println("4 and 6 connected? " + dsu.connected(4, 6)); // true

        // Merge the two big components
        dsu.union(3, 4);
        System.out.println("After union(3,4) — 1 and 6 connected? " + dsu.connected(1, 6)); // true
        System.out.println("Components: " + dsu.getComponentCount()); // 2
    }
}