Union by Rank + Path Compression

These two optimizations can be combined to achieve a $O(\alpha (n))$, where $\alpha$ is the inverse Ackermann function, which can be considered $O(1)$ on average.

Implementation

class UnionFind:
    def __init__(self, n: int) -> None:
        self.n = n
        self.parents = [i for i in range(self.n)]
        self.ranks = [1] * self.n

    def union(self, a: int, b: int) -> None:
        self._validate(a)
        self._validate(b)

        parent_a = self.find(a)
        parent_b = self.find(b)
        if parent_a == parent_b:
            return

        rank_a = self.ranks[a]
        rank_b = self.ranks[b]
        same_rank = rank_a == rank_b
        if rank_a <= rank_b:
            self.parents[parent_b] = parent_a
            if same_rank:
                self.ranks[b] += 1
        else:
            self.parents[parent_a] = parent_b

    def find(self, a: int) -> int:
        if a == self.parents[a]:
            return a

        self.parents[a] = self.find(self.parents[a])

        return self.parents[a]

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

    def _validate(self, a: int) -> None:
        assert 0 <= a < self.n

Again, applied to Find if Path Exists in Graph problem, our solution will be accepted.

from typing import List


class Solution:
    def validPath(
        self, n: int, edges: List[List[int]], source: int, destination: int
    ) -> bool:
        union_find = self.build_disjoint_set(n, edges)

        return union_find.connected(source, destination)

    def build_disjoint_set(self, n: int, edges: List[List[int]]) -> UnionFind:
        union_find = UnionFind(n)

        for a, b in edges:
            union_find.union(a, b)

        return union_find