# Union-Find (aussi appelé Disjoint Set Union) gère des ensembles disjoints avec 2 opérations :
# trouver(x) : identifier l'ensemble auquel x appartient
# union(x, y) : fusionner les ensembles de x et y
# Avec compression de chemin + union par rang : quasi O(1) amorti par opération (O(alpha(n))).

class UnionFind:
    def __init__(self, elements: list) -> None:
        self.parent = {e: e for e in elements}   # chaque élément est initialement son propre représentant
        self.rang = {e: 0 for e in elements}       # rang = hauteur approximative de l'arbre
        self.nb_composantes = len(elements)

    def trouver(self, x) -> "object":
        # Compression de chemin : rattache x DIRECTEMENT à la racine lors du parcours
        racine = x
        while self.parent[racine] != racine:
            racine = self.parent[racine]
        # deuxième passe : aplatit tout le chemin vers la racine trouvée
        while self.parent[x] != racine:
            self.parent[x], x = racine, self.parent[x]
        return racine

    def union(self, x, y) -> bool:
        racine_x, racine_y = self.trouver(x), self.trouver(y)
        if racine_x == racine_y:
            return False   # déjà connectés -- fusion inutile (indique un cycle si utilisé sur un graphe)

        # union par rang : le plus petit arbre est accroché SOUS le plus grand, garde l'arbre plat
        if self.rang[racine_x] < self.rang[racine_y]:
            racine_x, racine_y = racine_y, racine_x
        self.parent[racine_y] = racine_x
        if self.rang[racine_x] == self.rang[racine_y]:
            self.rang[racine_x] += 1

        self.nb_composantes -= 1
        return True

    def connectes(self, x, y) -> bool:
        return self.trouver(x) == self.trouver(y)

uf = UnionFind(["A", "B", "C", "D", "E"])
uf.union("A", "B")
uf.union("B", "C")
assert uf.connectes("A", "C") is True
assert uf.connectes("A", "D") is False
assert uf.nb_composantes == 3   # {A,B,C}, {D}, {E}

# --- Application : détecter un cycle dans un graphe non orienté ---
def graphe_contient_cycle(sommets: list, aretes: list[tuple]) -> bool:
    uf = UnionFind(sommets)
    for a, b in aretes:
        if not uf.union(a, b):   # union() retourne False si a et b étaient déjà connectés
            return True           # cette arête referme un cycle
    return False

assert graphe_contient_cycle(["A", "B", "C"], [("A", "B"), ("B", "C"), ("C", "A")]) is True
assert graphe_contient_cycle(["A", "B", "C"], [("A", "B"), ("B", "C")]) is False

# --- Application : compter les composantes connexes (ex: îles dans une grille, réseaux sociaux) ---
def nombre_iles(grille: list[list[int]]) -> int:
    lignes, colonnes = len(grille), len(grille[0])
    cellules_terre = [(i, j) for i in range(lignes) for j in range(colonnes) if grille[i][j] == 1]
    uf = UnionFind(cellules_terre)

    for i, j in cellules_terre:
        for di, dj in [(0, 1), (1, 0)]:   # ne vérifier que droite et bas évite les doublons
            ni, nj = i + di, j + dj
            if ni < lignes and nj < colonnes and grille[ni][nj] == 1:
                uf.union((i, j), (ni, nj))

    return uf.nb_composantes if cellules_terre else 0

grille = [
    [1, 1, 0, 0],
    [1, 0, 0, 1],
    [0, 0, 1, 1],
]
assert nombre_iles(grille) == 3

# --- Application : Kruskal (voir leçon MST) utilise directement cette structure ---
# --- Application : connectivité dynamique en réseau (ajout de connexions au fil du temps) ---
class ReseauDynamique:
    def __init__(self, machines: list[str]):
        self.uf = UnionFind(machines)

    def connecter(self, machine_a: str, machine_b: str) -> None:
        self.uf.union(machine_a, machine_b)

    def meme_reseau(self, machine_a: str, machine_b: str) -> bool:
        return self.uf.connectes(machine_a, machine_b)

reseau = ReseauDynamique(["srv1", "srv2", "srv3", "srv4"])
reseau.connecter("srv1", "srv2")
assert reseau.meme_reseau("srv1", "srv2") is True
assert reseau.meme_reseau("srv1", "srv3") is False