20. Cluster GPUs into Connected Training Groups

Problem

DeepMind's training infra periodically needs to know how many independent GPU groups exist in a slice of the interconnect fabric before deciding whether two experiment shards can share gradient all-reduce traffic.

Given n GPU nodes labeled 0..n-1 and a list of direct interconnect links between pairs of nodes, return the number of connected clusters ("training groups") the nodes partition into. Two nodes are in the same group if there is a path of links between them, directly or transitively.

Source: src/20_gpu_cluster_union_find.py

def count_training_groups(n: int, links: list[tuple[int, int]]) -> int:
    ...

Examples:
>>> count_training_groups(5, [(0, 1), (1, 2), (3, 4)])
2

>>> count_training_groups(4, [])
4

Step-by-Step Approach

1. Recognize "how many connected components does this graph have" as the signature use case for a union-find (disjoint set) structure, rather than building an adjacency list and running BFS/DFS from every unvisited node.
2. Initialize each of the n nodes as its own parent — n singleton groups to start.
3. For every link (a, b), find each node's current group root and, if they differ, union the two groups by attaching one root under the other.
4. Use path compression inside find so future lookups on the same nodes are near O(1) instead of walking a long chain.
5. Use union by rank (attach the shorter tree under the taller one) so trees stay shallow and no single union degrades to a linked list.
6. After processing every link, the number of distinct roots across all n nodes is the answer — find every node once more and count unique results.

The key insight is that union-find answers "are these two things connected" and "how many connected groups exist" in near-O(1) amortized time per operation, without ever materializing the full graph structure that BFS/DFS would need to traverse.

Reference solution

def count_training_groups(n: int, links: list[tuple[int, int]]) -> int:
    parent = list(range(n))
    rank = [0] * n

    def find(x: int) -> int:
        # path compression: point every visited node straight at the root
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: int, b: int) -> None:
        root_a, root_b = find(a), find(b)
        if root_a == root_b:
            return
        # union by rank: attach the shorter tree under the taller one
        if rank[root_a] < rank[root_b]:
            root_a, root_b = root_b, root_a
        parent[root_b] = root_a
        if rank[root_a] == rank[root_b]:
            rank[root_a] += 1

    for a, b in links:
        union(a, b)

    return len({find(i) for i in range(n)})

Key Functions & Tricks

• find(x) with path compression — flattens the tree on every lookup so later queries are near-instant
• union(a, b) by rank — keeps trees shallow by always attaching the shorter tree under the taller one
• len({find(i) for i in range(n)}) — one pass to collect and count distinct group roots
• Skip the union entirely when root_a == root_b — avoids a wasted no-op merge and (in some variants) helps detect a redundant/cycle-forming edge
• Amortized near-O(1) per operation with both optimizations together — effectively O(n + m·α(n)) for n nodes and m links

How to Recognize This Pattern