Union-find

Write a union-find class for the elements 0 to n - 1 (lessons 1 to 3). At the start every element is in a set of its own.

• find(x) returns the representative of x’s set: the same element for every member of the set.
• union(a, b) joins the two sets and returns True, or returns False if a and b were already in one set.
• connected(a, b) says whether they are in one set, size(x) is the number of elements in x’s set, and groups() is the number of sets.

d = DSU(5)
d.union(0, 1)        # True
d.union(1, 0)        # False: already together
d.union(1, 2)        # True
d.connected(0, 2)    # True
d.size(2)            # 3
d.groups()           # 3: {0, 1, 2}, {3}, {4}

Some tests join long chains. Write find with a loop, or make sure your trees stay short (lessons 2 and 3): a recursive find on a path of 100,000 nodes stops with RecursionError.

Limits

• Up to 200,000 elements and about a million operations. Use union by size or rank, and path compression: every operation should cost nearly O(1).