Implementation: Building a DSU Class

In this chapter we write a Python DSU class with both path compression and union by size, and explain it line by line. Then we port the same core routine to C++, Java and TypeScript. All four versions treat elements as integers from 0 to n-1.

Python implementation

class DSU:
    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.size = [1] * n
        self.count = n

    def find(self, x: int) -> int:
        root = x
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[x] != root:
            self.parent[x], x = root, self.parent[x]
        return root

    def union(self, a: int, b: int) -> bool:
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        self.count -= 1
        return True

    def connected(self, a: int, b: int) -> bool:
        return self.find(a) == self.find(b)

    def set_size(self, x: int) -> int:
        return self.size[self.find(x)]

Line by line

• self.parent = list(range(n)): every element points to itself, so we start with n singleton sets.
• self.size = [1] * n: the number of elements in each set. Only the value stored at a root is meaningful.
• self.count = n: the current number of sets, decremented on every successful merge.
• The first loop in find climbs parent links until it reaches the root.
• The second loop is path compression: it points each visited node at the root and moves on. The right-hand side is evaluated first, so self.parent[x], x = root, self.parent[x] safely advances to the original parent.
• Loops are used instead of recursion so that Python's default recursion limit (about 1000) is never hit.
• union finds both roots first. If they are equal, the elements are already together and it returns False; that return value is exactly what cycle detection needs.
• After making rb the smaller root, it hangs rb under ra and adds the sizes.
• connected and set_size are thin convenience wrappers around find.

Usage

dsu = DSU(7)
edges = [(0, 1), (1, 2), (3, 4), (5, 6), (2, 0)]
for a, b in edges:
    if not dsu.union(a, b):
        print("edge closes a cycle:", (a, b))   # (2, 0)

print(dsu.count)            # 3
print(dsu.connected(0, 2))  # True
print(dsu.set_size(4))      # 2

C++

#include <numeric>
#include <utility>
#include <vector>

struct DSU {
    std::vector<int> parent, size;
    int count;

    explicit DSU(int n) : parent(n), size(n, 1), count(n) {
        std::iota(parent.begin(), parent.end(), 0);
    }

    int find(int x) {
        int root = x;
        while (parent[root] != root) root = parent[root];
        while (parent[x] != root) {
            int next = parent[x];
            parent[x] = root;
            x = next;
        }
        return root;
    }

    bool unite(int a, int b) {
        int ra = find(a), rb = find(b);
        if (ra == rb) return false;
        if (size[ra] < size[rb]) std::swap(ra, rb);
        parent[rb] = ra;
        size[ra] += size[rb];
        --count;
        return true;
    }
};