Disjoint Set Union

A Disjoint Set Union (DSU)—also known as Union-Find—maintains a partition of $N$ elements (0..n) into disjoint connected components, supporting two operations in $\mathcal{O}(\alpha(n))$ nearly constant amortized time (where $\alpha(n) \le 4$ is the inverse Ackermann function for all practical $n$):

1. find(x): Returns the canonical root representative of x's component.
2. union(x, y): Merges the components containing x and y.

Path Compression + Union by Size (`O(alpha(N))` per query):

  Before `find(3)`:          After `find(3)` (Path Compression flattens tree!):
        (0)                              (0)
        /                              /  |  \
      (1)                            (1) (2) (3)
      /
    (2)
    /
  (3)

Why Combine Path Compression and Union by Size?

Complete Disjoint Set Union (UnionFind) Implementation in Rust

Tracking size: Vec<usize> alongside parent: Vec<usize> and num_components: usize lets you query the exact size of any element's connected component (component_size(x)) and the total number of remaining components in $O(1)$ time:

pub struct DisjointSetUnion {
    parent: Vec<usize>,
    size: Vec<usize>,
    num_components: usize,
}

impl DisjointSetUnion {
    pub fn new(n: usize) -> Self {
        Self {
            parent: (0..n).collect(),
            size: vec![1; n],
            num_components: n,
        }
    }

    /// Finds the canonical root of `x` with iterative path compression (path halving).
    #[inline]
    pub fn find(&mut self, mut x: usize) -> usize {
        while self.parent[x] != x {
            // Point `x` to its grandparent (`path halving`), flattening the tree
            // in a single pass with zero recursion overhead!
            let p = self.parent[x];
            self.parent[x] = self.parent[p];
            x = self.parent[x];
        }
        x
    }

    /// Merges the sets containing `a` and `b`.
    /// Returns `true` if `a` and `b` were in different sets (merged), or `false` if already connected.
    #[inline]
    pub fn union(&mut self, a: usize, b: usize) -> bool {
        let mut root_a = self.find(a);
        let mut root_b = self.find(b);
        if root_a == root_b {
            return false;
        }

        // Union by size: ensure `root_a` is the larger component
        if self.size[root_a] < self.size[root_b] {
            std::mem::swap(&mut root_a, &mut root_b);
        }
        self.parent[root_b] = root_a;
        self.size[root_a] += self.size[root_b];
        self.num_components -= 1;
        true
    }

    #[inline]
    pub fn same_set(&mut self, a: usize, b: usize) -> bool {
        self.find(a) == self.find(b)
    }

    #[inline]
    pub fn component_size(&mut self, x: usize) -> usize {
        let root = self.find(x);
        self.size[root]
    }

    #[inline]
    pub const fn components(&self) -> usize {
        self.num_components
    }
}

/// Kruskal's Minimum Spanning Tree (MST) algorithm in O(E log E) time using `DisjointSetUnion`.
pub fn kruskal_mst(num_vertices: usize, mut edges: Vec<(usize, usize, u64)>) -> Option<u64> {
    if num_vertices == 0 {
        return Some(0);
    }
    // Sort edges by weight ascending:
    edges.sort_unstable_by_key(|&(_u, _v, w)| w);

    let mut dsu = DisjointSetUnion::new(num_vertices);
    let mut total_weight = 0u64;

    for (u, v, weight) in edges {
        if dsu.union(u, v) {
            total_weight += weight;
            if dsu.components() == 1 {
                return Some(total_weight);
            }
        }
    }

    (dsu.components() == 1).then_some(total_weight)
}

fn main() {
    let edges = vec![
        (0, 1, 10),
        (0, 2, 6),
        (0, 3, 5),
        (1, 3, 15),
        (2, 3, 4),
    ];

    let mst_cost = kruskal_mst(4, edges).unwrap();
    // MST picks edges (2,3 w=4), (0,3 w=5), and (0,1 w=10) -> total = 19
    assert_eq!(mst_cost, 19);
    println!("Kruskal MST total weight: {mst_cost}");
}

Kruskal MST total weight: 19