Union Find Class

Implement a Disjoint Set Union (Union-Find) with path compression and union by size.

union-find disjoint-set data-structures

Implement a `UnionFind` class that maintains a disjoint set of `n` elements indexed from `0` to `n-1`. Initially, each element is its own parent and the size of each set is 1.

Your class must provide the following methods:

- `__init__(self, n: int)`: Initializes the structure with `n` elements. The elements are indexed from 0 to n-1. After initialization, each element is its own root and each set has size 1.
- `find(self, x: int) -> int`: Returns the root (representative) of the set containing `x`. This method must apply **path compression** along the path to the root (i.e., make the found root the direct parent of every node on the path).
- `union(self, x: int, y: int) -> None`: Merges the sets containing `x` and `y`. This method must use **union by size** (attach the root of the smaller set to the root of the larger set). If they are already in the same set, do nothing.
- `connected(self, x: int, y: int) -> bool`: Returns `True` if `x` and `y` are in the same set, `False` otherwise.

Input constraints: `1 <= n <= 10^4`, and test cases will not reference out-of-range indices.

Constraints

1 <= n <= 10^4. The number and type of method calls are unspecified but do not exceed typical limits. Your implementation should be efficient, ideally with amortized near-constant time using path compression and union by size.

Example

>>> uf = UnionFind(5)
>>> uf.union(0, 1)
>>> uf.connected(0, 1)
True
>>> uf.connected(0, 2)
False
>>> uf.find(1)
0
>>> uf.union(2, 3)
>>> uf.union(1, 2)
>>> uf.find(3)
0

10 points ~15 min

Recent Submissions

No submissions yet — hit Run Tests to try!

Hints

Store `parent` and `size` lists. Initially, parent[i] = i and size[i] = 1.

In `find`, recursively or iteratively update parent[x] to the root (path compression).

In `union`, use the roots from find. If different, attach the smaller tree to the larger tree and update size.

Test Results

Press Ctrl+Enter or click Run Tests to execute your code.