Union-Find (Disjoint Set) em Rust

Implemente Union-Find em Rust com union by rank e path compression. Diagramas de floresta, Kruskal MST e análise amortizada.

O Union-Find (também chamado de Disjoint Set Union, ou DSU) é uma estrutura de dados que gerencia uma coleção de conjuntos disjuntos (sem elementos em comum). Ela suporta duas operações fundamentais: Find (determinar a qual conjunto um elemento pertence) e Union (unir dois conjuntos em um só). Com as otimizações de union by rank e path compression, ambas as operações rodam em tempo amortizado quase constante: O(alfa(n)), onde alfa é a função inversa de Ackermann, que cresce tão lentamente que na prática e considerada constante.

Na prática, Union-Find e fundamental para algoritmos de grafos como o algoritmo de Kruskal para árvore geradora mínima, detecção de ciclos em grafos não direcionados, e problemas de conectividade dinâmica. Sua simplicidade de implementação e eficiência excepcional a tornam uma das estruturas mais elegantes da ciência da computação.

Como Funciona

A estrutura é representada como uma floresta (conjunto de árvores), onde cada elemento aponta para um pai. A raiz de cada árvore é o representante do conjunto.

Estado inicial: cada elemento é seu próprio conjunto
Elementos: {0, 1, 2, 3, 4, 5, 6, 7}

pai:   [0, 1, 2, 3, 4, 5, 6, 7]   (cada um aponta para si mesmo)
rank:  [0, 0, 0, 0, 0, 0, 0, 0]

Floresta:  0   1   2   3   4   5   6   7
           (8 árvores de 1 nó cada)

=== Union(1, 2) === "Junte os conjuntos de 1 e 2"

pai:   [0, 1, 1, 3, 4, 5, 6, 7]   (2 agora aponta para 1)
rank:  [0, 1, 0, 3, 4, 5, 6, 7]

Floresta:  0   1   3   4   5   6   7
               |
               2

=== Union(3, 4) ===

Floresta:  0   1   3   5   6   7
               |   |
               2   4

=== Union(1, 3) === "Junte {1,2} com {3,4}"

Como rank(1) = 1 e rank(3) = 1 → árvore de 3 fica sob 1
pai:   [0, 1, 1, 1, 3, 5, 6, 7]

Floresta:  0     1     5   6   7
              /  |
             2   3
                 |
                 4

=== Find(4) com Path Compression ===

Caminho: 4 → 3 → 1 (raiz)
Compressão: todos apontam direto para 1

pai:   [0, 1, 1, 1, 1, 5, 6, 7]   (4 agora aponta direto para 1)

Floresta:  0     1       5   6   7
            / | \
           2  3  4      (árvore achatada!)

A compressão de caminho achata a árvore durante o Find, fazendo todos os nós no caminho apontarem diretamente para a raiz. O union by rank sempre anexa a árvore menor à maior, mantendo as árvores rasas.

Complexidade

OperaçãoAmortizadoPior Caso (sem otim.)Espaço
FindO(alfa(n))O(n)O(n)
UnionO(alfa(n))O(n)O(n)
MakeSetO(1)O(1)O(1)
ConnectedO(alfa(n))O(n)O(1)
  • O(alfa(n)): a função inversa de Ackermann alfa(n) e <= 4 para qualquer valor prático de n (mesmo para n = 2^65536). Na prática, é constante.
  • Sem otimizações: Find pode ser O(n) se a árvore degenerar em lista (sem union by rank) e sem compressão de caminho.
  • Espaço O(n): dois arrays de tamanho n (pai e rank).

Implementação em Rust

/// Estrutura Union-Find com union by rank e path compression.
struct UnionFind {
    pai: Vec<usize>,
    rank: Vec<usize>,
    /// Número de conjuntos distintos.
    num_conjuntos: usize,
}

impl UnionFind {
    /// Cria uma estrutura com n elementos, cada um em seu próprio conjunto.
    fn novo(n: usize) -> Self {
        UnionFind {
            pai: (0..n).collect(),  // cada elemento é pai de si mesmo
            rank: vec![0; n],       // todas as árvores com rank 0
            num_conjuntos: n,
        }
    }

    /// Encontra o representante (raiz) do conjunto ao qual x pertence.
    /// Aplica compressão de caminho para achatar a árvore.
    fn find(&mut self, x: usize) -> usize {
        if self.pai[x] != x {
            // Compressão de caminho: faz x apontar direto para a raiz
            self.pai[x] = self.find(self.pai[x]);
        }
        self.pai[x]
    }

    /// Une os conjuntos que contêm x e y.
    /// Usa union by rank para manter as árvores rasas.
    /// Retorna true se os conjuntos eram distintos (e foram unidos).
    fn union(&mut self, x: usize, y: usize) -> bool {
        let raiz_x = self.find(x);
        let raiz_y = self.find(y);

        if raiz_x == raiz_y {
            return false; // Já estão no mesmo conjunto
        }

        // Union by rank: a árvore menor fica sob a maior
        match self.rank[raiz_x].cmp(&self.rank[raiz_y]) {
            std::cmp::Ordering::Less => {
                self.pai[raiz_x] = raiz_y;
            }
            std::cmp::Ordering::Greater => {
                self.pai[raiz_y] = raiz_x;
            }
            std::cmp::Ordering::Equal => {
                self.pai[raiz_y] = raiz_x;
                self.rank[raiz_x] += 1;
            }
        }

        self.num_conjuntos -= 1;
        true
    }

    /// Verifica se x e y estão no mesmo conjunto.
    fn conectados(&mut self, x: usize, y: usize) -> bool {
        self.find(x) == self.find(y)
    }

    /// Retorna o número de conjuntos distintos.
    fn conjuntos(&self) -> usize {
        self.num_conjuntos
    }

    /// Retorna o tamanho do conjunto ao qual x pertence.
    fn tamanho_conjunto(&mut self, x: usize) -> usize {
        let raiz = self.find(x);
        // Conta quantos elementos têm a mesma raiz
        (0..self.pai.len())
            .filter(|&i| {
                // Precisamos acessar self.pai sem mutabilidade aqui
                // Usamos uma busca direta sem compressão
                let mut atual = i;
                while self.pai[atual] != atual {
                    atual = self.pai[atual];
                }
                atual == raiz
            })
            .count()
    }
}

fn main() {
    let mut uf = UnionFind::novo(8);

    println!("Estado inicial: {} conjuntos", uf.conjuntos());

    // Unir elementos
    uf.union(1, 2);
    println!("Após union(1,2): {} conjuntos", uf.conjuntos());

    uf.union(3, 4);
    println!("Após union(3,4): {} conjuntos", uf.conjuntos());

    uf.union(1, 3);
    println!("Após union(1,3): {} conjuntos", uf.conjuntos());

    // Verificar conexões
    println!("\n2 e 4 conectados? {}", uf.conectados(2, 4)); // true
    println!("0 e 1 conectados? {}", uf.conectados(0, 1)); // false
    println!("1 e 4 conectados? {}", uf.conectados(1, 4)); // true

    // Unir mais
    uf.union(5, 6);
    uf.union(6, 7);
    uf.union(0, 5);
    println!("\nApós mais uniões: {} conjuntos", uf.conjuntos());
    println!("0 e 7 conectados? {}", uf.conectados(0, 7)); // true
}

Exemplo de Execução

Estado inicial: 8 conjuntos
Após union(1,2): 7 conjuntos
Após union(3,4): 6 conjuntos
Após union(1,3): 5 conjuntos

2 e 4 conectados? true
0 e 1 conectados? false
1 e 4 conectados? true

Após mais uniões: 2 conjuntos
0 e 7 conectados? true

Trace detalhado da compressão de caminho:

Após union(1,2), union(3,4), union(1,3):

Árvore:        1
              / \
             2   3
                 |
                 4

find(4):
  4 → pai[4]=3 → pai[3]=1 → pai[1]=1 (raiz!)
  Compressão: pai[4] = 1, pai[3] = 1

Árvore achatada:  1
                / | \
               2  3  4

Variações e Otimizações

1. Algoritmo de Kruskal para Árvore Geradora Mínima

O Union-Find é a peça-chave do algoritmo de Kruskal, que encontra a árvore geradora mínima de um grafo ponderado.

/// Aresta de um grafo ponderado.
#[derive(Debug, Clone)]
struct Aresta {
    origem: usize,
    destino: usize,
    peso: i64,
}

struct UnionFind {
    pai: Vec<usize>,
    rank: Vec<usize>,
}

impl UnionFind {
    fn novo(n: usize) -> Self {
        UnionFind {
            pai: (0..n).collect(),
            rank: vec![0; n],
        }
    }

    fn find(&mut self, x: usize) -> usize {
        if self.pai[x] != x {
            self.pai[x] = self.find(self.pai[x]);
        }
        self.pai[x]
    }

    fn union(&mut self, x: usize, y: usize) -> bool {
        let rx = self.find(x);
        let ry = self.find(y);
        if rx == ry { return false; }
        if self.rank[rx] < self.rank[ry] {
            self.pai[rx] = ry;
        } else if self.rank[rx] > self.rank[ry] {
            self.pai[ry] = rx;
        } else {
            self.pai[ry] = rx;
            self.rank[rx] += 1;
        }
        true
    }
}

/// Algoritmo de Kruskal: encontra a árvore geradora mínima.
/// Retorna as arestas da MST e o custo total.
fn kruskal(num_vertices: usize, mut arestas: Vec<Aresta>) -> (Vec<Aresta>, i64) {
    // Ordena arestas por peso crescente
    arestas.sort_by_key(|a| a.peso);

    let mut uf = UnionFind::novo(num_vertices);
    let mut mst = Vec::new();
    let mut custo_total = 0;

    for aresta in arestas {
        // Se a aresta conecta dois componentes diferentes, inclui na MST
        if uf.union(aresta.origem, aresta.destino) {
            custo_total += aresta.peso;
            mst.push(aresta);

            // MST tem exatamente V-1 arestas
            if mst.len() == num_vertices - 1 {
                break;
            }
        }
    }

    (mst, custo_total)
}

fn main() {
    // Grafo exemplo:
    //     0 ---4--- 1
    //     |  \      |
    //     8   2     6
    //     |     \   |
    //     3 ---3--- 2
    //      \       /
    //       7    9
    //        \  /
    //         4

    let arestas = vec![
        Aresta { origem: 0, destino: 1, peso: 4 },
        Aresta { origem: 0, destino: 3, peso: 8 },
        Aresta { origem: 0, destino: 2, peso: 2 },
        Aresta { origem: 1, destino: 2, peso: 6 },
        Aresta { origem: 2, destino: 3, peso: 3 },
        Aresta { origem: 2, destino: 4, peso: 9 },
        Aresta { origem: 3, destino: 4, peso: 7 },
    ];

    let (mst, custo) = kruskal(5, arestas);

    println!("Árvore Geradora Mínima (Kruskal):");
    for aresta in &mst {
        println!("  {} -- {} (peso: {})", aresta.origem, aresta.destino, aresta.peso);
    }
    println!("Custo total: {}", custo);
}

2. Detecção de Ciclos em Grafos Não Direcionados

struct UnionFind {
    pai: Vec<usize>,
    rank: Vec<usize>,
}

impl UnionFind {
    fn novo(n: usize) -> Self {
        UnionFind {
            pai: (0..n).collect(),
            rank: vec![0; n],
        }
    }

    fn find(&mut self, x: usize) -> usize {
        if self.pai[x] != x {
            self.pai[x] = self.find(self.pai[x]);
        }
        self.pai[x]
    }

    fn union(&mut self, x: usize, y: usize) -> bool {
        let rx = self.find(x);
        let ry = self.find(y);
        if rx == ry { return false; }
        if self.rank[rx] < self.rank[ry] {
            self.pai[rx] = ry;
        } else {
            self.pai[ry] = rx;
            if self.rank[rx] == self.rank[ry] {
                self.rank[rx] += 1;
            }
        }
        true
    }
}

/// Detecta se um grafo não direcionado contém ciclos.
fn tem_ciclo(num_vertices: usize, arestas: &[(usize, usize)]) -> bool {
    let mut uf = UnionFind::novo(num_vertices);

    for &(u, v) in arestas {
        if !uf.union(u, v) {
            // u e v já estão no mesmo componente → aresta cria ciclo!
            println!("  Ciclo detectado: aresta ({}, {})", u, v);
            return true;
        }
    }
    false
}

fn main() {
    // Grafo sem ciclo (árvore)
    let arestas1 = vec![(0, 1), (1, 2), (2, 3)];
    println!("Grafo 1 (sem ciclo):");
    println!("  Tem ciclo? {}\n", tem_ciclo(4, &arestas1));

    // Grafo com ciclo: 0-1-2-0
    let arestas2 = vec![(0, 1), (1, 2), (2, 0)];
    println!("Grafo 2 (com ciclo):");
    println!("  Tem ciclo? {}", tem_ciclo(3, &arestas2));
}

3. Union-Find com Union by Size (alternativa ao rank)

/// Union-Find com union by size — permite consultar o tamanho de cada conjunto em O(1).
struct UnionFindSize {
    pai: Vec<usize>,
    tamanho: Vec<usize>, // tamanho do conjunto (válido apenas para raízes)
}

impl UnionFindSize {
    fn novo(n: usize) -> Self {
        UnionFindSize {
            pai: (0..n).collect(),
            tamanho: vec![1; n],
        }
    }

    fn find(&mut self, x: usize) -> usize {
        if self.pai[x] != x {
            self.pai[x] = self.find(self.pai[x]);
        }
        self.pai[x]
    }

    fn union(&mut self, x: usize, y: usize) -> bool {
        let rx = self.find(x);
        let ry = self.find(y);
        if rx == ry { return false; }

        // Anexa o menor conjunto ao maior
        if self.tamanho[rx] < self.tamanho[ry] {
            self.pai[rx] = ry;
            self.tamanho[ry] += self.tamanho[rx];
        } else {
            self.pai[ry] = rx;
            self.tamanho[rx] += self.tamanho[ry];
        }
        true
    }

    /// Retorna o tamanho do conjunto de x em O(alfa(n)).
    fn tamanho_de(&mut self, x: usize) -> usize {
        let raiz = self.find(x);
        self.tamanho[raiz]
    }
}

fn main() {
    let mut uf = UnionFindSize::novo(10);

    uf.union(0, 1);
    uf.union(2, 3);
    uf.union(0, 2);

    println!("Tamanho do conjunto de 0: {}", uf.tamanho_de(0)); // 4
    println!("Tamanho do conjunto de 5: {}", uf.tamanho_de(5)); // 1

    uf.union(5, 6);
    uf.union(6, 7);
    uf.union(0, 5);
    println!("Tamanho do conjunto de 7: {}", uf.tamanho_de(7)); // 7
}

Aplicações Práticas

  • Árvore Geradora Mínima (Kruskal): o Union-Find é peça central do algoritmo de Kruskal, determinando se uma aresta conecta dois componentes distintos.
  • Componentes conexos: em grafos dinâmicos onde arestas são adicionadas incrementalmente, Union-Find mantém os componentes conexos de forma eficiente.
  • Detecção de ciclos: em grafos não direcionados, uma aresta que conecta dois vértices já no mesmo conjunto cria um ciclo.
  • Equivalência de pixels em imagens: em processamento de imagens, Union-Find agrupa pixels adjacentes com características similares (segmentação).
  • Redes sociais: determinar se duas pessoas pertencem ao mesmo grupo ou comunidade.
  • Percolação: simulações físicas de percolação usam Union-Find para detectar quando um caminho se forma entre extremidades.

Comparação com Alternativas

CaracterísticaUnion-FindBFS/DFSMatriz de adjacência
Verificar conexãoO(alfa(n))O(V + E)O(1)
Unir componentesO(alfa(n))-O(V)
Listar componenteO(n)O(V + E)O(V)
EspaçoO(n)O(V + E)O(V^2)
Suporta remoçãoNãoSimSim
Dinâmico (inserção)SimRecomputarSim
ImplementaçãoSimplesMédiaSimples

Union-Find é ideal quando as operações são apenas adições (não suporta remoção de arestas eficientemente). Para problemas com remoções, BFS/DFS é mais adequado.

Exercícios Práticos

  1. Número de ilhas: Dada uma grade m x n de ‘0’s (água) e ‘1’s (terra), conte o número de ilhas. Uma ilha é cercada por água e formada conectando terras adjacentes horizontalmente ou verticalmente. Use Union-Find para resolver.

  2. Contas redundantes: Em uma rede social, pares (email1, email2) indicam que pertencem à mesma pessoa. Dado uma lista de pares, determine quantas pessoas distintas existem.

  3. Grafo bipartido: Usando Union-Find, verifique se um grafo não direcionado é bipartido. Dica: para cada vértice, seus vizinhos devem estar no mesmo conjunto (oposto ao do vértice).

  4. Union-Find com rollback: Implemente uma versão do Union-Find que suporta desfazer a última operação de union (sem path compression, usando union by rank com histórico). Isso é útil em problemas offline de divisão e conquista.

Veja Também