Algoritmo Union Find

Abaixo está uma implementação simples, idiomática e correta de Union Find (Disjoint Set Union – DSU) em Python, usando path compression e union by rank.

Implementação

class UnionFind:
    def __init__(self, n: int):
        # parent[i] = pai do elemento i
        self.parent = list(range(n))
        # rank[i] = estimativa da altura da árvore
        self.rank = [0] * n
 
    def find(self, x: int) -> int:
        # Path compression
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
 
    def union(self, x: int, y: int) -> bool:
        root_x = self.find(x)
        root_y = self.find(y)
 
        if root_x == root_y:
            return False  # já estão no mesmo conjunto
 
        # Union by rank
        if self.rank[root_x] < self.rank[root_y]:
            self.parent[root_x] = root_y
        elif self.rank[root_x] > self.rank[root_y]:
            self.parent[root_y] = root_x
        else:
            self.parent[root_y] = root_x
            self.rank[root_x] += 1
 
        return True
 
    def connected(self, x: int, y: int) -> bool:
        return self.find(x) == self.find(y)

Exemplo de uso

uf = UnionFind(5)
 
uf.union(0, 1)
uf.union(1, 2)
 
print(uf.connected(0, 2))  # True
print(uf.connected(0, 3))  # False
 
uf.union(3, 4)
uf.union(2, 4)
 
print(uf.connected(0, 4))  # True

Propriedades importantes

• find: amortizado O(α(n))
• union: amortizado O(α(n))
• α(n) é a função inversa de Ackermann, cresce mais lentamente que log n
• Na prática: quase O(1)

Observação conceitual importante

• Rank não é atualizado no path compression
  • Rank é apenas uma estimativa histórica da altura
  • Path compression pode “achatar” a árvore sem quebrar a validade do rank