Union-find in Rust

It is the first week of college. Asha becomes friends with Chen. Later, Chen becomes friends with Dia. Nobody introduced Asha to Dia, but they are now in the same circle, through Chen.

Two questions keep coming up. Are these two people in the same circle? And when two people become friends, how do you merge their circles into one?

That is exactly what union-find does. It keeps track of items split into groups, where every item is in exactly one group. Mathematicians call these disjoint sets, and disjoint only means "with no overlap". It has two moves:

• find tells you which group an item is in
• union merges two groups into one

Reach for it when things join up over time and you keep asking "are these connected?": friend circles, networks that gain links, regions in an image, or building the cheapest set of roads that joins every town.

Here is the idea, in the plainest form. Every person points at one other person. A group's leader is the one who points at themselves. To find out which group you are in, follow the pointers upward until you reach the leader. Two people are in the same group when they reach the same leader.

To merge two groups, take one leader and make them point at the other leader. That is all. Two circles are now one.

Try it on four people, numbered 0 to 3. Merge 0 with 1, so 0 points at 1. Merge 1 with 2, so 1 points at 2. Merge 2 with 3, so 2 points at 3. You get a chain: 0 points at 1, which points at 2, which points at 3. To find 0's leader, you walk three steps. With a thousand people merged this way, it would be a thousand steps. The simple version can get slow.

Two small ideas fix this, and they work together.

Path compression. When a find walks up a chain, it learns who the leader is. So on the way back, it points every person it passed straight at the leader. In the picture, one find(3) turns the long chain into a flat group, and every later find takes one step. The walk is done once, and then it never has to be done again.

Union by rank. Rank is a number for each leader that stands for how tall their tree could be. When merging, hang the shorter tree under the taller one, never the other way. Then a merge cannot make a tall tree much taller. The rank only goes up by one when two trees of the same height are joined.

Used together, they make find and union so fast that people call the cost "almost constant". The exact name is the inverse Ackermann function, written α(n), and it grows so slowly that it never goes above 4 for any amount of data that could fit in the universe.

Remember: point everyone you pass at the leader, and hang the shorter tree under the taller one.

First, the simple version with no tricks. Each element starts as its own leader.

/// The simplest version. `parent[i]` is the person `i` points to. A person who
/// points at themselves is the leader of their group.
pub struct SlowUnionFind {
    parent: Vec<usize>,
}

impl SlowUnionFind {
    /// `n` people, each in a group of one.
    pub fn new(n: usize) -> Self {
        SlowUnionFind {
            parent: (0..n).collect(),
        }
    }

    /// Follow the pointers up until you reach a leader.
    pub fn find(&self, mut x: usize) -> usize {
        while self.parent[x] != x {
            x = self.parent[x];
        }
        x
    }

    /// Joins two groups by hanging one leader under the other.
    pub fn union(&mut self, a: usize, b: usize) {
        let (leader_a, leader_b) = (self.find(a), self.find(b));
        self.parent[leader_a] = leader_b;
    }

    /// How many steps `x` is from its leader.
    pub fn depth(&self, mut x: usize) -> usize {
        let mut steps = 0;
        while self.parent[x] != x {
            x = self.parent[x];
            steps += 1;
        }
        steps
    }
}

parent[i] is who i points at. new makes everyone point at themselves, so everyone starts alone. find follows pointers until it finds someone who points at themselves. union finds both leaders and hangs one under the other. depth is only here so you can see the problem.

let mut friends = SlowUnionFind::new(4);
friends.union(0, 1);
friends.union(1, 2);
friends.union(2, 3);

// Every join hung the newest leader under the next one: a long chain.
println!("{}", friends.depth(0));
assert_eq!(friends.depth(0), 3);
assert_eq!(friends.find(0), friends.find(3));

Three merges built a chain, and element 0 is three steps from its leader, exactly as in the example. Now the fast version.

/// Union-find with the two speed-ups: path compression and union by rank.
pub struct UnionFind {
    parent: Vec<usize>,
    /// A cap on how tall the tree under each leader can be.
    rank: Vec<u8>,
    groups: usize,
}

impl UnionFind {
    /// `n` elements, each in a group of one.
    pub fn new(n: usize) -> Self {
        UnionFind {
            parent: (0..n).collect(),
            rank: vec![0; n],
            groups: n,
        }
    }

    /// How many separate groups there are right now.
    pub fn groups(&self) -> usize {
        self.groups
    }

    /// The leader of `x`'s group. Along the way, everyone passed is pointed
    /// straight at the leader, so the next search is short.
    pub fn find(&mut self, x: usize) -> usize {
        let mut leader = x;
        while self.parent[leader] != leader {
            leader = self.parent[leader];
        }
        // Second walk: path compression.
        let mut current = x;
        while current != leader {
            let next = self.parent[current];
            self.parent[current] = leader;
            current = next;
        }
        leader
    }

    /// Merges the groups of `a` and `b`. Returns `false` if they were already
    /// the same group.
    pub fn union(&mut self, a: usize, b: usize) -> bool {
        let (mut top, mut bottom) = (self.find(a), self.find(b));
        if top == bottom {
            return false;
        }
        // Hang the shorter tree under the taller one, so trees stay short.
        if self.rank[top] < self.rank[bottom] {
            std::mem::swap(&mut top, &mut bottom);
        }
        self.parent[bottom] = top;
        if self.rank[top] == self.rank[bottom] {
            self.rank[top] += 1; // two equal trees joined make a taller one
        }
        self.groups -= 1;
        true
    }

    pub fn connected(&mut self, a: usize, b: usize) -> bool {
        self.find(a) == self.find(b)
    }
}

Go through it in order.

• rank holds the height cap for each leader, and groups counts the groups. Both are kept up to date as things merge.
• find walks up twice. The first walk finds the leader. The second walk goes up again from x, and points each person at the leader. That second walk is path compression.
• union finds the two leaders. If they are the same, the two were already together, so it returns false and changes nothing. Otherwise std::mem::swap makes sure top is the leader with the larger rank, and bottom is hung underneath. Only if the ranks were equal does the winner's rank go up by one. Every real merge also takes one off groups.
• connected is just two finds compared.

let mut friends = UnionFind::new(6);
friends.union(0, 1);
friends.union(2, 3);
assert!(!friends.connected(1, 2));

friends.union(1, 2); // joins {0, 1} with {2, 3}
assert!(friends.connected(0, 3));
assert!(!friends.connected(0, 4));
assert!(!friends.union(0, 3)); // already friends, so nothing changes

println!("{}", friends.groups());
assert_eq!(friends.groups(), 3); // {0,1,2,3}, {4}, {5}

Four friends end up in one group, and 4 and 5 stay alone, so three groups remain. The Playground link on the code above runs this.

The standard library has no union-find. (The petgraph crate has one, if you ever need it in a project.) What Rust does give you are the pieces: a Vec for the pointers, std::mem::swap for the rank step, and a HashMap for something real data needs.

Real data has names, not numbers. Union-find is built on Vec positions, so you turn each name into a number first, the same way as the station names in lesson 10.

use std::collections::HashMap;

// Union-find works on numbers, so give each name a number first.
let names = ["Asha", "Ben", "Chen", "Dia"];
let mut number = HashMap::new();
for (i, name) in names.iter().enumerate() {
    number.insert(*name, i);
}

let mut friends = UnionFind::new(names.len());
friends.union(number["Asha"], number["Chen"]);
friends.union(number["Chen"], number["Dia"]);

println!("{}", friends.connected(number["Asha"], number["Dia"]));
assert!(friends.connected(number["Asha"], number["Dia"]));
assert!(!friends.connected(number["Asha"], number["Ben"]));

The structure only ever sees numbers. The table is how you talk to it in words.

Remember: union-find works on numbers. Keep a HashMap beside it to turn names into numbers.

Without the two tricks, a find can walk a chain of n items, so it is O(n). With both, α(n) is at most 4 in practice, so you can read it as constant. One thing union-find cannot do is split a group back apart. It only merges.

These use the UnionFind you just built. The Playground link on each problem includes it. Each starter has a todo!(), a placeholder that panics if it runs. Replace it with your code.

pub fn count_groups(n: usize, pairs: &[(usize, usize)]) -> usize {
    todo!("start with n groups; every real merge removes one")
}

pub fn count_groups(n: usize, pairs: &[(usize, usize)]) -> usize {
    let mut uf = UnionFind::new(n);
    for &(a, b) in pairs {
        uf.union(a, b);
    }
    uf.groups()
}

pub fn first_redundant(n: usize, edges: &[(usize, usize)]) -> Option<(usize, usize)> {
    todo!("the first edge that joins two things already joined")
}

pub fn first_redundant(n: usize, edges: &[(usize, usize)]) -> Option<(usize, usize)> {
    let mut uf = UnionFind::new(n);
    for &(a, b) in edges {
        // `union` says false when both ends were already joined, which means
        // this edge closes a loop.
        if !uf.union(a, b) {
            return Some((a, b));
        }
    }
    None
}

pub fn min_cost_connect(n: usize, roads: &[(usize, usize, u32)]) -> Option<u32> {
    todo!("the cheapest set of roads that joins everything")
}

pub fn min_cost_connect(n: usize, roads: &[(usize, usize, u32)]) -> Option<u32> {
    // Cheapest first. Skip any road whose ends are already joined.
    let mut sorted = roads.to_vec();
    sorted.sort_by_key(|&(_, _, cost)| cost);

    let mut uf = UnionFind::new(n);
    let mut total = 0;
    for (a, b, cost) in sorted {
        if uf.union(a, b) {
            total += cost;
        }
    }
    // One group left means every place is reachable.
    (uf.groups() <= 1).then_some(total)
}

Union-find is a small structure with a big reputation. A few habits help.

• Spot the signals. "Are these connected?", "how many groups?", "merge", "same component", "does this edge make a cycle?", and "cheapest way to connect everything" all point here.
• Name both tricks. Saying "path compression and union by rank" shows you know the real structure. Even one of them gets you most of the speed, but name both.
• Say the cost the right way. "Nearly constant, formally O(α(n)), where α is the inverse Ackermann function and never exceeds 4 in practice" is the answer an interviewer wants.
• Compare it with graph search. For a fixed graph, BFS or DFS from lesson 10 can count the pieces too. Union-find wins when edges arrive over time and you keep asking "are these connected right now?".
• Know what it cannot do. It merges groups but cannot split them. If a problem deletes links, you need another idea.
• Use the return value. union returning false for an edge that was already redundant is a neat, short way to detect a cycle.
• In Rust, explain the &mut. Mentioning that find needs &mut self because it compresses paths, and why that is worth it, shows you understand ownership as well as the algorithm.

Made by VeraStack Labs