Union-Find class with component count

Implement a class UnionFind(n) over the elements 0 .. n-1, which start out in n separate sets:

• find(x) -> int: the representative (root) of x's set. Two elements are in the same set exactly when their find values are equal.
• union(a, b) -> bool: merge the sets of a and b. Return True if they were different sets, False if they were already together (then nothing changes).
• connected(a, b) -> bool: whether a and b are in the same set.
• count: an attribute holding the current number of sets.

uf = UnionFind(5)
uf.union(0, 1)      # True: merged two sets
uf.union(1, 0)      # False: already in the same set
uf.connected(0, 1)  # True
uf.connected(0, 2)  # False
uf.count            # 4
uf.find(1) == uf.find(0)   # True

Constraints: 0 <= n <= 2 * 10^5; the tests do hundreds of thousands of operations, including one long chain.

Use path compression (point every node you walk past straight at the root) and union by rank or size (hang the shorter tree under the taller one). Without them a chain of unions makes find O(n) and the large test times out. If you skip union by rank, write find as a loop: a recursive one would hit Python's recursion limit on a 200,000-long chain before compression kicks in.