Union-Find in Python, Java and TypeScript

A parent array with both optimisations, counting components and detecting a cycle.

None of Python, Java and TypeScript ships a union-find in its standard library, so it is written by hand, in about twenty lines. This card covers a complete implementation with union by rank and path compression in all three languages, and the two jobs it is most often given: counting connected components and detecting a cycle.

In union-find each item stores a parent, and an item that is its own parent is the root of its group. find follows parents to the root, and union links one root under the other. Union by rank links the root of the shorter tree under the taller one, using a rank that bounds each tree's height, and path compression makes find point every item it visits at the root.

Each version keeps two arrays, parent and rank, and a count of components, the number of separate groups, which starts at the number of items. Each listing joins six items with the edges 0-1, 2-3, 1-3, 0-2 and 4-5 and reports the count after each edge.

The same union-find in each language

union_find_code.pypy

class UnionFind:
    def __init__(self, size):
        self.parent = list(range(size))
        self.rank = [0] * size
        self.components = size

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a, b):
        root_a, root_b = self.find(a), self.find(b)
        if root_a == root_b:
            return False
        if self.rank[root_a] > self.rank[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_a] = root_b
        if self.rank[root_a] == self.rank[root_b]:
            self.rank[root_b] += 1
        self.components -= 1
        return True


groups = UnionFind(6)
for a, b in [(0, 1), (2, 3), (1, 3), (0, 2), (4, 5)]:
    if groups.union(a, b):
        print(f"union({a}, {b}): joined, {groups.components} components")
    else:
        print(f"union({a}, {b}): already joined, so edge {a}-{b} closes a cycle")
print(f"0 and 3 connected: {groups.find(0) == groups.find(3)}")
print(f"0 and 4 connected: {groups.find(0) == groups.find(4)}")
print(f"components: {groups.components}")

OutputRan with python:3.14.8-slim-trixie (Python 3.14.8)

union(0, 1): joined, 5 components
union(2, 3): joined, 4 components
union(1, 3): joined, 3 components
union(0, 2): already joined, so edge 0-2 closes a cycle
union(4, 5): joined, 2 components
0 and 3 connected: True
0 and 4 connected: False
components: 2

UnionFindCode.javajava

public class UnionFindCode {
    static class UnionFind {
        private final int[] parent;
        private final int[] rank;
        private int components;

        UnionFind(int size) {
            parent = new int[size];
            rank = new int[size];
            components = size;
            for (int i = 0; i < size; i++) {
                parent[i] = i;
            }
        }

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

        boolean union(int a, int b) {
            int rootA = find(a);
            int rootB = find(b);
            if (rootA == rootB) {
                return false;
            }
            if (rank[rootA] > rank[rootB]) {
                int swap = rootA;
                rootA = rootB;
                rootB = swap;
            }
            parent[rootA] = rootB;
            if (rank[rootA] == rank[rootB]) {
                rank[rootB]++;
            }
            components--;
            return true;
        }

        int components() {
            return components;
        }
    }

    public static void main(String[] args) {
        UnionFind groups = new UnionFind(6);
        int[][] edges = {{0, 1}, {2, 3}, {1, 3}, {0, 2}, {4, 5}};
        for (int[] edge : edges) {
            int a = edge[0];
            int b = edge[1];
            if (groups.union(a, b)) {
                System.out.println("union(" + a + ", " + b + "): joined, " + groups.components() + " components");
            } else {
                System.out.println("union(" + a + ", " + b + "): already joined, so edge " + a + "-" + b + " closes a cycle");
            }
        }
        System.out.println("0 and 3 connected: " + (groups.find(0) == groups.find(3)));
        System.out.println("0 and 4 connected: " + (groups.find(0) == groups.find(4)));
        System.out.println("components: " + groups.components());
    }
}

OutputRan with eclipse-temurin:25.0.4.1_1-jdk-noble (Java 25.0.4.1)

union(0, 1): joined, 5 components
union(2, 3): joined, 4 components
union(1, 3): joined, 3 components
union(0, 2): already joined, so edge 0-2 closes a cycle
union(4, 5): joined, 2 components
0 and 3 connected: true
0 and 4 connected: false
components: 2

union_find_code.tsts

class UnionFind {
  private readonly parent: number[];
  private readonly rank: number[];
  components: number;

  constructor(size: number) {
    this.parent = Array.from({ length: size }, (_, i) => i);
    this.rank = new Array<number>(size).fill(0);
    this.components = size;
  }

  find(x: number): number {
    if (this.parent[x] !== x) {
      this.parent[x] = this.find(this.parent[x]);
    }
    return this.parent[x];
  }

  union(a: number, b: number): boolean {
    let rootA = this.find(a);
    let rootB = this.find(b);
    if (rootA === rootB) {
      return false;
    }
    if (this.rank[rootA] > this.rank[rootB]) {
      [rootA, rootB] = [rootB, rootA];
    }
    this.parent[rootA] = rootB;
    if (this.rank[rootA] === this.rank[rootB]) {
      this.rank[rootB]++;
    }
    this.components--;
    return true;
  }
}

const groups = new UnionFind(6);
const edges: ReadonlyArray<readonly [number, number]> = [[0, 1], [2, 3], [1, 3], [0, 2], [4, 5]];
for (const [a, b] of edges) {
  if (groups.union(a, b)) {
    console.log(`union(${a}, ${b}): joined, ${groups.components} components`);
  } else {
    console.log(`union(${a}, ${b}): already joined, so edge ${a}-${b} closes a cycle`);
  }
}
console.log(`0 and 3 connected: ${groups.find(0) === groups.find(3)}`);
console.log(`0 and 4 connected: ${groups.find(0) === groups.find(4)}`);
console.log(`components: ${groups.components}`);

OutputRan with node:24.22.0-trixie-slim (Node 24.22.0), TypeScript 5.9.3

union(0, 1): joined, 5 components
union(2, 3): joined, 4 components
union(1, 3): joined, 3 components
union(0, 2): already joined, so edge 0-2 closes a cycle
union(4, 5): joined, 2 components
0 and 3 connected: true
0 and 4 connected: false
components: 2

The highlighted lines carry the ideas. find is recursive, and the assignment on the way back out is path compression: every item on the path is pointed at the root the recursion returned. union returns false when both roots are the same, links root under root otherwise, and lowers the count by one each time it succeeds.

Those two details do the two jobs. Each successful union merges two components, so the count after the last edge is the number of connected components, groups of items joined by some chain of edges: 2 here, {0, 1, 2, 3} and {4, 5}. A union that finds both ends already under one root has met an edge between two items that were already connected, so adding it closes a cycle, a path that returns to where it started. In the run, edge 0-2 did that, because 0-1, 1-3 and 3-2 already linked them.

The recursive find goes as deep as the tree, which union by rank keeps within log₂ n, so the recursion stays shallow; a version without rank can recurse thousands of levels on a long chain and hit Python's recursion limit. Java's new int[size] starts every rank at 0, as it should, while the TypeScript version fills its rank array explicitly.

← → move through sections and cards