# Generated from the published Markdown; run scripts/build_language_examples.py to refresh.
# Independent examples may reuse top-level names. Python remains the explanation language.


# topics/union-find/parent-forests.md
def find_root(parent, x):
    while parent[x] != x:
        x = parent[x]
    return x

# topics/union-find/union-compression.md
class DisjointSet:
    def __init__(self, n):
        if type(n) is not int or n < 0:
            raise ValueError("n must be a nonnegative integer")
        self.parent = list(range(n))
        self.size = [1] * n
        self.components = n

    def _check(self, x):
        if type(x) is not int or not 0 <= x < len(self.parent):
            raise IndexError("member index out of range")

    def find(self, x):
        self._check(x)
        root = x
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[x] != x:
            next_x = self.parent[x]
            self.parent[x] = root
            x = next_x
        return root

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        self.components -= 1
        return True

    def connected(self, a, b):
        return self.find(a) == self.find(b)

    def __len__(self):
        return len(self.parent)

    def is_empty(self):
        return len(self) == 0

    def component_count(self):
        return self.components

    def component_size(self, x):
        return self.size[self.find(x)]

# topics/union-find/union-find-applications.md
def connection_forest(n, edges):
    sets = DisjointSet(n)
    chosen, rejected = [], []
    for a, b in edges:
        if sets.union(a, b):
            chosen.append((a, b))
        else:
            rejected.append((a, b))
    return chosen, rejected, sets.components
