Quick Union

For this kind of implementation, we store the parent vertexes of each node.

For example:

Would have the following values:

index             :  0 1 2 3 4 5
value (root node) : [0 0 1 3 3 4]

Implementation

The implementation would look something like:

class UnionFind:
    def __init__(self, n: int) -> None:
        self.n = n
        self.parents = [i for i in range(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:
            self.parents[parent_b] = parent_a

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

        return 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

Usage

Compared to the Quick Find, we have passed more test cases in Find if Path Exists in Graph, but this is not enough. We are still getting a Time Limit Exceeded.

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