Learn Zig Series (#122) - Union-Find

scipio

Published: 29 Jul 2026 › Updated: 29 Jul 2026 Learn Zig Series (#122) - Union-Find

Learn Zig Series (#122) - Union-Find

What will I learn?

What the union-find (or disjoint-set) structure actually answers -- the deceptively simple question "are these two things in the same group yet?" -- and why a full graph traversal is the wrong tool for it;
The three implementations in order of getting-smarter: quick-find, quick-union, and the real deal with union by size and path compression;
Why those two little optimizations drag the cost from linear down to inverse Ackermann -- a function so flat it is a constant for any input that will ever exist;
How to key a disjoint-set by arbitrary names (strings, not just array indices) using a std.StringHashMap, and how Zig's error set keeps that honest;
A property test that pins the fast version to a dead-simple reference, reusing the two-implementations-one-oracle trick from last episode;
The single most famous use of the whole structure -- Kruskal's minimum spanning tree -- where the boolean we return from unite turns out to be exactly "is this edge in the tree?";
How the same structure looks in C, Rust and Go, and why mutating during a read (path compression) is the load-bearing surprise in every one of them.

Requirements

A working modern computer running macOS, Windows or Ubuntu;
An installed Zig 0.14+ distribution (download from ziglang.org);
Episode 121 (Topological Sort) still fresh -- we settle its three exercises in full compilable code at the very top, and the cycle-reporting solution reuses the gray-vertex stack we built there almost verbatim;
The std.ArrayList and allocator discipline from episodes 7 and 22, plus the std.mem.sort and std.StringHashMap comfort from the collections episodes;
The ambition to learn Zig programming.

Difficulty

Advanced

Curriculum (of the `Learn Zig Series`):

Learn Zig Series (#122) - Union-Find

At the end of last episode I promised you a structure "so small you could write it on a napkin and so fast its running time is, for all practical purposes, a flat constant per operation". Well -- here it is. Topological sort answered "in what full order do I process everything?" by knowing all the edges up front and asking once. Today's structure answers a smaller, sneakier question that real systems ask millions of times as they run: are these two things in the same group yet? Are these two pixels part of the same blob? Are these two accounts already in the same cluster? Would adding this one edge to my growing spanning tree close a loop, or bridge two islands? Re-running a graph traversal for every one of those queries would be madness. The structure that makes each answer near-instant is called union-find (or, more formally, the disjoint-set data structure), and it is one of those rare things that feels like cheating the first time you meet it. But first -- the debt. Last episode's three topological-sort exercises, in full compilable code ;-)

Solutions to Episode 121 Exercises

Exercise 1 -- report the actual cycle, not just its existence. Last episode's topoSortDfs returned a bare error.Cycle and threw away the single most useful piece of information: which vertices are stuck in the loop. But we already had that information in our hands at the moment of detection -- the gray vertices currently on the stack are exactly the active path, and when a back-edge points into a gray vertex w, the cycle is the run of the stack from the frame whose vertex is w up to the top. So I return a tagged union -- either an order or a cycle -- and on detection I scan the stack for w, copy that tail out, and hand it back. Note the memory choreography on the cycle branch: I allocate the cycle slice with try before tearing down the half-built order list, so if that allocation fails the errdefer order.deinit(gpa) still fires exactly once (no double free), and if it succeeds I free order by hand and return:

const std = @import("std");

const Graph = struct {
    gpa: std.mem.Allocator,
    adj: std.ArrayList(std.ArrayList(usize)),
    directed: bool,

    fn init(gpa: std.mem.Allocator, n: usize, directed: bool) !Graph {
        var adj: std.ArrayList(std.ArrayList(usize)) = .empty;
        try adj.ensureTotalCapacity(gpa, n);
        var i: usize = 0;
        while (i < n) : (i += 1) adj.appendAssumeCapacity(.empty);
        return .{ .gpa = gpa, .adj = adj, .directed = directed };
    }
    fn deinit(self: *Graph) void {
        for (self.adj.items) |*list| list.deinit(self.gpa);
        self.adj.deinit(self.gpa);
    }
    fn addEdge(self: *Graph, from: usize, to: usize) !void {
        try self.adj.items[from].append(self.gpa, to);
        if (!self.directed) try self.adj.items[to].append(self.gpa, from);
    }
    fn neighbors(self: *const Graph, v: usize) []const usize {
        return self.adj.items[v].items;
    }
};

const TopoResult = union(enum) {
    order: []usize,
    cycle: []usize,
};

fn topoSortReport(gpa: std.mem.Allocator, g: *const Graph) !TopoResult {
    const n = g.adj.items.len;
    const Color = enum(u8) { white, gray, black };
    const color = try gpa.alloc(Color, n);
    defer gpa.free(color);
    @memset(color, .white);

    var order: std.ArrayList(usize) = .empty;
    errdefer order.deinit(gpa);

    const Frame = struct { v: usize, idx: usize };
    var stack: std.ArrayList(Frame) = .empty;
    defer stack.deinit(gpa);

    for (0..n) |start| {
        if (color[start] != .white) continue;
        try stack.append(gpa, .{ .v = start, .idx = 0 });
        color[start] = .gray;
        while (stack.items.len > 0) {
            const top = &stack.items[stack.items.len - 1];
            const nbs = g.neighbors(top.v);
            if (top.idx < nbs.len) {
                const w = nbs[top.idx];
                top.idx += 1;
                switch (color[w]) {
                    .gray => {
                        var start_idx: usize = 0;
                        for (stack.items, 0..) |f, i| {
                            if (f.v == w) {
                                start_idx = i;
                                break;
                            }
                        }
                        const cyc = try gpa.alloc(usize, stack.items.len - start_idx);
                        for (stack.items[start_idx..], 0..) |f, i| cyc[i] = f.v;
                        order.deinit(gpa);
                        return .{ .cycle = cyc };
                    },
                    .white => {
                        color[w] = .gray;
                        try stack.append(gpa, .{ .v = w, .idx = 0 });
                    },
                    .black => {},
                }
            } else {
                color[top.v] = .black;
                try order.append(gpa, top.v);
                stack.items.len -= 1;
            }
        }
    }
    std.mem.reverse(usize, order.items);
    return .{ .order = try order.toOwnedSlice(gpa) };
}

test "cycle reporting recovers the vertices of the cycle" {
    const gpa = std.testing.allocator;
    var g = try Graph.init(gpa, 3, true);
    defer g.deinit();
    try g.addEdge(0, 1);
    try g.addEdge(1, 2);
    try g.addEdge(2, 0);
    const result = try topoSortReport(gpa, &g);
    switch (result) {
        .order => |o| {
            defer gpa.free(o);
            try std.testing.expect(false);
        },
        .cycle => |c| {
            defer gpa.free(c);
            try std.testing.expectEqual(@as(usize, 3), c.len);
            var seen = [_]bool{false} ** 3;
            for (c) |v| seen[v] = true;
            try std.testing.expect(seen[0] and seen[1] and seen[2]);
        },
    }
}

The tagged union is the honest Zig way to say "this returns one of two structurally different things", and the caller's switch is forced by the compiler to handle both arms. On the 0 -> 1 -> 2 -> 0 graph we get {0, 1, 2} in some rotation -- exactly the message a build tool would print as "circular dependency: 0 -> 1 -> 2 -> 0".

Exercise 2 -- enumerate all topological orders by backtracking. One valid order is what a build system needs; all of them is what shows you, concretely, how much freedom the DAG really has. The recipe is classic backtracking on Kahn's in-degree idea: at each step, every currently-ready vertex (in-degree zero, not yet used) is a legal next choice, so we try each one -- mark it used, decrement its neighbours' in-degrees, recurse, then undo all three mutations before trying the next candidate. When the current path reaches length n we snapshot it with dupe. I cap it at 8 vertices with an assert, because the count of valid orders explodes combinatorially and you do not want to discover that with a 30-node graph ;-)

const Backtracker = struct {
    gpa: std.mem.Allocator,
    g: *const Graph,
    n: usize,
    indeg: []usize,
    used: []bool,
    current: std.ArrayList(usize),
    results: std.ArrayList([]usize),

    fn recurse(self: *Backtracker) !void {
        if (self.current.items.len == self.n) {
            const snap = try self.gpa.dupe(usize, self.current.items);
            errdefer self.gpa.free(snap);
            try self.results.append(self.gpa, snap);
            return;
        }
        for (0..self.n) |v| {
            if (self.used[v] or self.indeg[v] != 0) continue;
            self.used[v] = true;
            for (self.g.neighbors(v)) |w| self.indeg[w] -= 1;
            try self.current.append(self.gpa, v);
            try self.recurse();
            _ = self.current.pop();
            for (self.g.neighbors(v)) |w| self.indeg[w] += 1;
            self.used[v] = false;
        }
    }
};

fn allTopoOrders(gpa: std.mem.Allocator, g: *const Graph) !std.ArrayList([]usize) {
    const n = g.adj.items.len;
    std.debug.assert(n <= 8);
    const indeg = try gpa.alloc(usize, n);
    defer gpa.free(indeg);
    @memset(indeg, 0);
    for (0..n) |u| {
        for (g.neighbors(u)) |v| indeg[v] += 1;
    }
    const used = try gpa.alloc(bool, n);
    defer gpa.free(used);
    @memset(used, false);

    var bt = Backtracker{
        .gpa = gpa,
        .g = g,
        .n = n,
        .indeg = indeg,
        .used = used,
        .current = .empty,
        .results = .empty,
    };
    defer bt.current.deinit(gpa);
    errdefer {
        for (bt.results.items) |r| gpa.free(r);
        bt.results.deinit(gpa);
    }
    try bt.recurse();
    return bt.results;
}

The undo step is the whole soul of backtracking: every mutation on the way down has a mirror-image on the way back up, so the shared indeg/used scratch always returns to precisely the state the caller left it in. I keep the recursion state in a little Backtracker struct so the recursive method mutates the same arrays across the whole descent in stead of threading half a dozen parameters through every call. And notice the ownership split in allTopoOrders: the scratch (indeg, used, current) is freed here via defer, but the results list is handed back to the caller, who now owns every snapshot slice inside it:

test "backtracking enumerates every valid topological order" {
    const gpa = std.testing.allocator;
    var g = try Graph.init(gpa, 4, true);
    defer g.deinit();
    try g.addEdge(0, 2); // 0 before 2
    try g.addEdge(1, 2); // 1 before 2, and 3 is free
    var orders = try allTopoOrders(gpa, &g);
    defer {
        for (orders.items) |o| gpa.free(o);
        orders.deinit(gpa);
    }
    try std.testing.expectEqual(@as(usize, 8), orders.items.len);
    for (orders.items) |o| {
        try std.testing.expect(try isValidTopoOrder(gpa, &g, o));
    }
}

Two constraints (0 -> 2, 1 -> 2) on four vertices, with vertex 3 completely free, admit exactly 8 distinct valid orders -- and every single one passes the isValidTopoOrder checker from last episode. That is the "freedom" of a DAG made countable.

Exercise 3 -- the longest dependency chain (the critical path). In project scheduling this is the chain of tasks that decides your minimum possible finish time: no matter how much you parallelise, you cannot beat the longest must-happen-in-sequence run. On a DAG it is embarrassingly cheap: walk the vertices in topological order, and for each edge u -> v relax longest[v] = max(longest[v], longest[u] + 1). Because we visit in topo order, by the time we reach any vertex all of its predecessors have already been finalised -- the exact same "predecessors first" guarantee that made the order useful in the first place:

fn topoSortKahn(gpa: std.mem.Allocator, g: *const Graph) !?[]usize {
    const n = g.adj.items.len;
    const indeg = try gpa.alloc(usize, n);
    defer gpa.free(indeg);
    @memset(indeg, 0);
    for (0..n) |u| {
        for (g.neighbors(u)) |v| indeg[v] += 1;
    }
    var queue: std.ArrayList(usize) = .empty;
    defer queue.deinit(gpa);
    for (0..n) |v| {
        if (indeg[v] == 0) try queue.append(gpa, v);
    }
    var order: std.ArrayList(usize) = .empty;
    errdefer order.deinit(gpa);
    var head: usize = 0;
    while (head < queue.items.len) {
        const u = queue.items[head];
        head += 1;
        try order.append(gpa, u);
        for (g.neighbors(u)) |v| {
            indeg[v] -= 1;
            if (indeg[v] == 0) try queue.append(gpa, v);
        }
    }
    if (order.items.len != n) {
        order.deinit(gpa);
        return null;
    }
    return try order.toOwnedSlice(gpa);
}

fn longestChain(gpa: std.mem.Allocator, g: *const Graph) !?usize {
    const order = (try topoSortKahn(gpa, g)) orelse return null;
    defer gpa.free(order);
    const n = g.adj.items.len;
    const longest = try gpa.alloc(usize, n);
    defer gpa.free(longest);
    for (longest) |*x| x.* = 1;
    var best: usize = 0;
    for (order) |u| {
        if (longest[u] > best) best = longest[u];
        for (g.neighbors(u)) |v| {
            if (longest[u] + 1 > longest[v]) longest[v] = longest[u] + 1;
        }
    }
    return best;
}

test "longest dependency chain is the critical path length" {
    const gpa = std.testing.allocator;
    var g = try Graph.init(gpa, 4, true);
    defer g.deinit();
    try g.addEdge(0, 1); // libc -> libm
    try g.addEdge(0, 2); // libc -> app
    try g.addEdge(1, 2); // libm -> app
    try g.addEdge(2, 3); // app  -> tests
    const chain = (try longestChain(gpa, &g)).?;
    try std.testing.expectEqual(@as(usize, 4), chain);
}

On the little libc -> libm -> app -> tests build graph the answer is 4 -- the whole four-node chain is the critical path, because every step genuinely waits on the one before it. If you had a fifth library that app also needed but which didn't depend on libm, it would not lengthen the chain; it would just be another job you can slot in alongside. Debt paid ;-) Now -- groups.

The question union-find is built to answer

Here is the setup. You have n elements -- number them 0 to n-1 -- and they start out as n separate singleton groups, each element alone in its own club. Over time, two operations happen, in any interleaving:

union(a, b): merge the group containing a and the group containing b into one.
find(a) / connected(a, b): which group is a in? Are a and b in the same group?

That is the entire interface. It sounds almost too plain to deserve a name, but look at where it shows up: flood-filling connected regions in an image, detecting whether adding an edge to a graph creates a cycle, grouping equivalent variables in a type checker, the percolation problem in physics, and -- the crown jewel we will end on -- building a minimum spanning tree. In every one of those, the same two questions get asked over and over as the structure grows: merge these two, and are these two together yet?

You could answer connected with a BFS every time, but that is O(V + E) per query, and we are going to ask it a staggering number of times. Union-find gets each query down to (essentially) constant time by never traversing the graph at all. It stores just enough structure to answer "same group?" and nothing more.

The naive baseline: quick-find

The most obvious representation: an array id where id[x] is the group label of element x. Two elements are in the same group exactly when they carry the same label. connected is then a single comparison -- gloriously O(1). The catch is union: to merge two groups you must sweep the entire array and relabel every member of one group to match the other, which is O(n) per union:

const std = @import("std");

const QuickFind = struct {
    gpa: std.mem.Allocator,
    id: []usize,

    fn init(gpa: std.mem.Allocator, n: usize) !QuickFind {
        const id = try gpa.alloc(usize, n);
        for (id, 0..) |*x, i| x.* = i;
        return .{ .gpa = gpa, .id = id };
    }
    fn deinit(self: *QuickFind) void {
        self.gpa.free(self.id);
    }
    fn connected(self: *const QuickFind, a: usize, b: usize) bool {
        return self.id[a] == self.id[b];
    }
    fn unite(self: *QuickFind, a: usize, b: usize) void {
        const pa = self.id[a];
        const pb = self.id[b];
        if (pa == pb) return;
        for (self.id) |*x| {
            if (x.* == pa) x.* = pb;
        }
    }
};

test "quick-find groups elements correctly" {
    const gpa = std.testing.allocator;
    var qf = try QuickFind.init(gpa, 10);
    defer qf.deinit();
    qf.unite(0, 1);
    qf.unite(1, 2);
    qf.unite(8, 9);
    try std.testing.expect(qf.connected(0, 2));
    try std.testing.expect(qf.connected(8, 9));
    try std.testing.expect(!qf.connected(2, 8));
}

unite is a plain method, not !void -- notice it cannot fail, because it allocates nothing (we grabbed the whole id array once, in init). That is a small but real Zig tell: a function's signature advertises whether it can fail, and quick-find's unite honestly cannot. It is correct, it is dead simple, and for a workload of mostly-queries it is even fast. But n unions cost O(n^2), and once n gets into the millions that quadratic is a wall. We keep quick-find around, though, because its obvious correctness makes it the perfect reference oracle later.

The forest idea: quick-union

Flip the trade-off. In stead of a flat label per element, give each element a parent pointer. Follow parents upward and you eventually hit a root (an element that is its own parent); that root is the group's identity. Elements in the same tree share the same root. Now union is cheap -- just point one root at the other, O(1) after you find them -- but find has to walk up the tree:

const QuickUnion = struct {
    gpa: std.mem.Allocator,
    parent: []usize,

    fn init(gpa: std.mem.Allocator, n: usize) !QuickUnion {
        const parent = try gpa.alloc(usize, n);
        for (parent, 0..) |*p, i| p.* = i;
        return .{ .gpa = gpa, .parent = parent };
    }
    fn deinit(self: *QuickUnion) void {
        self.gpa.free(self.parent);
    }
    fn find(self: *const QuickUnion, x: usize) usize {
        var root = x;
        while (self.parent[root] != root) root = self.parent[root];
        return root;
    }
    fn connected(self: *const QuickUnion, a: usize, b: usize) bool {
        return self.find(a) == self.find(b);
    }
    fn unite(self: *QuickUnion, a: usize, b: usize) void {
        self.parent[self.find(a)] = self.find(b);
    }
};

test "quick-union links roots" {
    const gpa = std.testing.allocator;
    var qu = try QuickUnion.init(gpa, 6);
    defer qu.deinit();
    qu.unite(0, 1);
    qu.unite(2, 3);
    qu.unite(1, 3);
    try std.testing.expect(qu.connected(0, 2));
    try std.testing.expect(!qu.connected(0, 4));
}

We swapped the expensive operation. The trouble is that naive quick-union can build degenerate trees: unite in a pathological order -- 0 under 1, 1 under 2, 2 under 3, ... -- and you get a single tall stick, at which point find walks O(n) up the chain and we have gained nothing. Having said that, we are one insight away from fixing that permanently, and it is a good one.

The two optimizations that change the complexity class

Two independent ideas turn quick-union from "sometimes a linked list" into "flatter than you can measure". They compose, and together they are what people mean when they say union-find.

Union by size (its sibling is union by rank): when merging two trees, always hang the smaller tree under the larger tree's root, never the other way around. A tree can only get taller when it is merged with one at least as big as itself, which caps its height at O(log n) even before the second trick. We track each root's subtree size in a size array and compare before linking.

Path compression: find already has to walk from a node up to its root -- so, while we are up there anyway, point every node we passed directly at the root. The next find on any of them is then a single hop. The path we walked to answer this query flattens the tree for all future queries. It is the closest thing to a free lunch in this whole series.

Put both together and you get the canonical disjoint-set. I also carry a count of how many distinct sets remain (it starts at n and drops by one on every successful merge -- handy for "how many connected components?"), and I make unite return a bool: true if it actually merged two different sets, false if they were already together. Hold onto that boolean -- it is about to do real work in Kruskal's algorithm:

const DisjointSet = struct {
    gpa: std.mem.Allocator,
    parent: []usize,
    size: []usize,
    count: usize,

    fn init(gpa: std.mem.Allocator, n: usize) !DisjointSet {
        const parent = try gpa.alloc(usize, n);
        errdefer gpa.free(parent);
        const size = try gpa.alloc(usize, n);
        for (parent, 0..) |*p, i| p.* = i;
        @memset(size, 1);
        return .{ .gpa = gpa, .parent = parent, .size = size, .count = n };
    }
    fn deinit(self: *DisjointSet) void {
        self.gpa.free(self.parent);
        self.gpa.free(self.size);
    }
    fn find(self: *DisjointSet, x: usize) usize {
        var root = x;
        while (self.parent[root] != root) root = self.parent[root];
        var cur = x;
        while (self.parent[cur] != root) {
            const next = self.parent[cur];
            self.parent[cur] = root; // compress: point straight at the root
            cur = next;
        }
        return root;
    }
    fn connected(self: *DisjointSet, a: usize, b: usize) bool {
        return self.find(a) == self.find(b);
    }
    fn unite(self: *DisjointSet, a: usize, b: usize) bool {
        const ra = self.find(a);
        const rb = self.find(b);
        if (ra == rb) return false; // already in the same set
        if (self.size[ra] < self.size[rb]) {
            self.parent[ra] = rb;
            self.size[rb] += self.size[ra];
        } else {
            self.parent[rb] = ra;
            self.size[ra] += self.size[rb];
        }
        self.count -= 1;
        return true;
    }
    fn setSize(self: *DisjointSet, x: usize) usize {
        return self.size[self.find(x)];
    }
};

There is a subtle, telling detail here: find takes *DisjointSet, not *const DisjointSet. It mutates the structure -- that is path compression rewriting parent pointers -- even though the caller thinks of it as a read-only "which group?" query. Consequently connected also cannot be const. Zig makes you spell that out, and it is honest: a query that secretly rewrites memory is exactly the kind of thing a language should force into the open in stead of hiding. Here is the behaviour, including the count bookkeeping and the "already connected" boolean:

test "disjoint set tracks components and sizes" {
    const gpa = std.testing.allocator;
    var ds = try DisjointSet.init(gpa, 10);
    defer ds.deinit();
    try std.testing.expectEqual(@as(usize, 10), ds.count);
    try std.testing.expect(ds.unite(0, 1));
    try std.testing.expect(ds.unite(1, 2));
    try std.testing.expect(ds.unite(3, 4));
    try std.testing.expect(!ds.unite(0, 2)); // already connected -> false
    try std.testing.expectEqual(@as(usize, 7), ds.count);
    try std.testing.expect(ds.connected(0, 2));
    try std.testing.expect(!ds.connected(2, 3));
    try std.testing.expectEqual(@as(usize, 3), ds.setSize(0));
}

Ten elements, three successful merges, one rejected merge -- and count lands on 7, setSize(0) reports the three-element club {0, 1, 2}. The structure knows, at all times and in constant space, exactly how the groups stand.

Why this is (almost) constant time

Here is the punchline that makes complexity theorists slightly emotional. With both union by size and path compression, a sequence of m operations on n elements runs in O(m * alpha(n)), where alpha is the inverse Ackermann function. The Ackermann function grows so violently fast that its inverse grows unfathomably slowly: alpha(n) is at most 4 for any n up to the number of atoms in the observable universe, and beyond. So while it is not technically O(1), it is a constant for every input that will ever be constructed by anything, anywhere. Each operation is, for all practical purposes, flat. That is the "feels like cheating" I promised.

If you want compression without the two-pass walk in find, there is a lovely one-pass variant called path halving: as you climb, point each node at its grandparent, skipping every other link. It flattens almost as well, touches each pointer once, and is friendlier to the cache:

fn findHalving(self: *DisjointSet, x: usize) usize {
    var cur = x;
    while (self.parent[cur] != cur) {
        self.parent[cur] = self.parent[self.parent[cur]]; // halve the path
        cur = self.parent[cur];
    }
    return cur;
}

Whether the halving variant or the full two-pass compression wins is genuinely workload-dependent and worth measuring rather than guessing -- exactly the profiling discipline from episode 34. Both land in the same near-constant complexity class; the difference is constant factors and cache behaviour.

Keying by name, not index (and where the errors live)

Array-indexed elements are clean for teaching, but real inputs come as names: file paths, package identifiers, account handles. The standard move is a mapping layer -- a std.StringHashMap from name to a dense integer id, allocated on first sight, with the union-find running underneath on those ids. This is also where Zig's error handling earns its keep: idOf might have to grow three different containers, so it is !usize, and everything that builds on it (unite, connected) becomes fallible too and rides the normal try machinery:

const NamedSets = struct {
    gpa: std.mem.Allocator,
    parent: std.ArrayList(usize),
    size: std.ArrayList(usize),
    index: std.StringHashMap(usize),

    fn init(gpa: std.mem.Allocator) NamedSets {
        return .{
            .gpa = gpa,
            .parent = .empty,
            .size = .empty,
            .index = std.StringHashMap(usize).init(gpa),
        };
    }
    fn deinit(self: *NamedSets) void {
        self.parent.deinit(self.gpa);
        self.size.deinit(self.gpa);
        self.index.deinit();
    }
    fn idOf(self: *NamedSets, name: []const u8) !usize {
        if (self.index.get(name)) |i| return i;
        const i = self.parent.items.len;
        try self.parent.append(self.gpa, i);
        try self.size.append(self.gpa, 1);
        try self.index.put(name, i);
        return i;
    }
    fn findIndex(self: *NamedSets, x: usize) usize {
        var root = x;
        while (self.parent.items[root] != root) root = self.parent.items[root];
        var cur = x;
        while (self.parent.items[cur] != root) {
            const next = self.parent.items[cur];
            self.parent.items[cur] = root;
            cur = next;
        }
        return root;
    }
    fn unite(self: *NamedSets, a: []const u8, b: []const u8) !bool {
        const ra = self.findIndex(try self.idOf(a));
        const rb = self.findIndex(try self.idOf(b));
        if (ra == rb) return false;
        if (self.size.items[ra] < self.size.items[rb]) {
            self.parent.items[ra] = rb;
            self.size.items[rb] += self.size.items[ra];
        } else {
            self.parent.items[rb] = ra;
            self.size.items[ra] += self.size.items[rb];
        }
        return true;
    }
    fn connected(self: *NamedSets, a: []const u8, b: []const u8) !bool {
        const ra = self.findIndex(try self.idOf(a));
        const rb = self.findIndex(try self.idOf(b));
        return ra == rb;
    }
};

test "named union-find groups by string key" {
    const gpa = std.testing.allocator;
    var ns = NamedSets.init(gpa);
    defer ns.deinit();
    _ = try ns.unite("libc", "libm");
    _ = try ns.unite("libm", "app");
    try std.testing.expect(try ns.connected("libc", "app"));
    try std.testing.expect(!(try ns.connected("libc", "shell")));
}

The one sharp edge to keep in mind (and I say this so you don't get bitten): std.StringHashMap stores the key slices, not copies, so the strings you hand it must outlive the map. String literals like "libc" live for the whole program, so the test is fine, but if your names came from a buffer you were about to reuse, you would need to dupe them and free them in deinit. Zig does not paper over that lifetime question -- it is yours to answer, honestly, which is the whole philosophy.

Testing: the fast version against the obvious one

We built quick-find precisely because it is obviously correct -- so it makes a perfect oracle for the clever DisjointSet. This is the same "two implementations, one property" shape from last episode: throw hundreds of random union sequences at both, then assert they agree about every pair. If a subtle bug ever crept into the path-compression logic, this test would catch the disagreement immediately, and because it all runs on std.testing.allocator, any leak fails the round too:

test "property: optimized DSU agrees with the naive quick-find baseline" {
    const gpa = std.testing.allocator;
    var prng = std.Random.DefaultPrng.init(0x5AFE);
    const rand = prng.random();
    var round: usize = 0;
    while (round < 200) : (round += 1) {
        const n = rand.intRangeAtMost(usize, 1, 30);
        var qf = try QuickFind.init(gpa, n);
        defer qf.deinit();
        var ds = try DisjointSet.init(gpa, n);
        defer ds.deinit();
        const m = rand.intRangeAtMost(usize, 0, 40);
        for (0..m) |_| {
            const a = rand.uintLessThan(usize, n);
            const b = rand.uintLessThan(usize, n);
            qf.unite(a, b);
            _ = ds.unite(a, b);
        }
        for (0..n) |a| {
            for (0..n) |b| {
                try std.testing.expectEqual(qf.connected(a, b), ds.connected(a, b));
            }
        }
    }
}

Two hundred random scenarios, up to 30 elements and 40 merges each, and the two implementations never disagree on a single pair. The slow-but-obvious version is the perfect chaperone for the fast-but-clever one -- when you cannot easily prove the clever code correct, pin it to code you can trust by eye.

The killer application: Kruskal's minimum spanning tree

Now the payoff that put union-find on the map. A minimum spanning tree (MST) of a weighted, connected graph is the cheapest set of edges that keeps everything in one piece -- think laying cable to connect every building for the least total length. Kruskal's algorithm is almost insultingly simple: sort all edges cheapest-first, then walk them in that order and greedily keep each edge -- but only if it connects two currently-separate components. If an edge's two endpoints are already in the same component, adding it would just close a redundant loop, so we skip it.

And "are these two endpoints already in the same component?" is precisely the question union-find answers in near-constant time. Even sweeter: our unite already returns a bool that is true exactly when the two were in different sets. So the entire "should this edge be in the tree?" decision is the return value of unite. No separate connected check needed -- unite tests and merges in one shot:

const WEdge = struct { u: usize, v: usize, w: u32 };

fn kruskalMst(gpa: std.mem.Allocator, n: usize, edges: []WEdge) !u64 {
    std.mem.sort(WEdge, edges, {}, struct {
        fn lt(_: void, a: WEdge, b: WEdge) bool {
            return a.w < b.w;
        }
    }.lt);
    var ds = try DisjointSet.init(gpa, n);
    defer ds.deinit();
    var total: u64 = 0;
    for (edges) |e| {
        if (ds.unite(e.u, e.v)) total += e.w; // joins two components -> keep it
    }
    return total;
}

test "kruskal builds a minimum spanning tree via union-find" {
    const gpa = std.testing.allocator;
    var edges = [_]WEdge{
        .{ .u = 0, .v = 1, .w = 1 },
        .{ .u = 1, .v = 2, .w = 2 },
        .{ .u = 0, .v = 2, .w = 4 }, // redundant + heavier -> skipped
        .{ .u = 2, .v = 3, .w = 3 },
    };
    const cost = try kruskalMst(gpa, 4, &edges);
    try std.testing.expectEqual(@as(u64, 6), cost); // 1 + 2 + 3
}

Four vertices, four candidate edges. Sorted by weight we take 0-1 (weight 1) and 1-2 (weight 2), then reach 0-2 (weight 4) and reject it because 0 and 2 are already in one component -- that unite returns false, so it is skipped -- then take 2-3 (weight 3). Total 1 + 2 + 3 = 6, and the redundant heavier edge never makes the cut. The whole correctness of Kruskal rests on that forementioned boolean from unite. Bam, jonguh ;-) Sort the edges, let union-find veto the loops, and the greedy choice is provably optimal.

How C, Rust, and Go do it

In C, it is two int arrays (parent and size or rank) and two short functions, and honestly the C version is nearly as terse as ours -- union-find is one of the few advanced structures that is genuinely pleasant in C. The classic gotcha is writing find recursively for the path-compression elegance and then blowing the call stack on an adversarial input; the iterative two-pass or the halving loop sidesteps that, which is why we wrote it iteratively from the start. No deinit, no allocator threading, but also no one watching for the free you forgot.

In Rust, there is no disjoint-set in std, though petgraph ships a UnionFind and the crate ecosystem has several. Rolling your own runs parent and rank as Vec<usize>, and here the borrow checker delivers a genuinely instructive shock: path compression mutates during find, so find must take &mut self, which means you cannot hold a shared borrow of the structure while calling it. That is the same *DisjointSet-not-*const truth we hit in Zig, except Rust refuses to compile the alternative rather than merely letting you write const and regret it. Same lesson, louder enforcement.

In Go, it is parent and size as []int slices, the garbage collector spares you the cleanup, and generics (post-1.18) finally let you write a name-keyed version over comparable keys without interface{} gymnastics. sort.Slice handles Kruskal's edge sort in one closure, much like our std.mem.sort. And, this being Go, someone will propose sharding the structure across goroutines for parallelism -- which runs headfirst into the fact that path compression is a write on every find, so concurrent union-find needs real care (atomics, or lock-free tricks) and is a research topic unto itself, not a casual go keyword. Four languages, one tiny structure, and the same load-bearing surprise in every one: the innocent-looking find is a mutation. That is the thing to remember ;-)

Exercises

Union by rank in stead of size. Rewrite DisjointSet to track a rank (an upper bound on tree height) rather than a subtree size: attach the lower-rank root under the higher-rank one, and only bump a rank when two equal-rank trees merge. Then reuse the property test from this episode -- swap the rank version in for ds and confirm it still agrees with QuickFind on 200 random scenarios. The point is that either heuristic gives the near-constant bound; convince yourself of that experimentally.
A rollback-able disjoint set. Some algorithms (offline dynamic connectivity, DSU-on-tree) need to unite and then undo the last union. Path compression makes that impossible to reverse cheaply, so build a variant that uses union-by-rank without compression, pushes a little record of each merge onto an undo stack, and exposes rollback() that pops the last merge and restores both the parent pointer and the rank it changed. Test that unite, unite, rollback leaves the structure identical to having done just the first unite.
Counting islands. Given a grid of booleans (true = land), map each cell (r, c) to the id r * cols + c, union every land cell with its orthogonal land neighbours (right and down are enough on a single pass), and report the number of distinct land components using the count field. Test it on a small hand-drawn grid where you can eyeball the answer -- say a 4x4 with three separated blobs -- and confirm you get 3. This is the connected-components problem that image processing calls "blob labelling", solved in a dozen lines.

Where this is heading

We now have a structure that merges groups and answers "together yet?" in near-constant time, keyed by index or by name, and we have watched it collapse the whole of Kruskal's MST into a single greedy loop. Notice what union-find is built to do, though: it only ever grows. You can merge groups forever, but there is no split, no forget, no remove -- it accumulates and never lets go. Plenty of real structures face the opposite pressure. Picture a cache sitting in front of a slow database or a network call: it has a hard size limit, and every time it fills up it must throw something away to make room -- and the thing to evict is the one you have not touched in the longest while. That demands answering, in O(1), both "give me this key's value" and "which entry is the least recently used?" -- two questions that pull toward two different data structures, which we will have to marry into one. That marriage of a hash map and a doubly linked list is where we go next ;-)

Thanks for sticking with the whole graph arc -- we covered quit some ground. De groeten, and see you in the next one!

@scipio

Leave Learn Zig Series (#122) - Union-Find to:

Written by scipio

Does it matter who's right, or who's left?

Best Posts From scipio

We have not curated any of scipio's posts yet. But you can encourage our curation team to review posts by visiting them regularly and by referring other readers. Because we give priority to frequently read content.

Learn Zig Series (#122) - Union-Find