union_find.py

__all__ = ["find", "union", "is_group"]


def find(node: int, parents: dict[int, int]) -> int:
    if parents[node] != node:
        parents[node] = find(parents[node], parents=parents)
    return parents[node]


def union(n1: int, n2: int, parents: dict[int, int]) -> None:
    n1 = find(n1, parents=parents)
    n2 = find(n2, parents=parents)
    if n1 < n2:
        parents[n2] = n1
    else:
        parents[n1] = n2


def is_group(n1: int, n2: int, parents: dict[int, int]) -> bool:
    return find(n1, parents=parents) == find(n2, parents=parents)


if __name__ == "__main__":
    n_nodes = 10
    parents = {node: node for node in range(1, n_nodes + 1)}

    union(1, 2, parents)
    union(2, 3, parents)
    union(3, 4, parents)
    union(4, 5, parents)

    union(6, 7, parents)
    union(7, 8, parents)

    print(f"1, 5 same group?: {is_group(1, 3, parents)}")
    print(f"3, 7 same group?: {is_group(3, 7, parents)}")