Implementando uma Árvore Balanceada com Treap

A estrutura de dados em questão deve suportar seis operações: inserção, remoção, consulta de valor pelo ranking, consulta de ranking por valor, consulta de antecessor e consulta de sucessor.

Para as operações de consulta, uma Árvore de Busca Binária (BST) é uma solução adequada. Para suportar todas as consultas, cada nó na BST precisa armazenar três informações:

  1. key: O valor armazenado no nó.
  2. cnt: A contagem de ocorrências desse valor no nó (para lidar com duplicatas).
  3. size: O número total de nós na subárvore enraizada neste nó (incluindo ele mesmo e suas duplicatas).

Consulta de Valor pelo Ranking

Dado um ranking rank, a função getKey(u, rank) busca o valor correspondente na subárvore enraizada em u.

  • Se o rank for menor ou igual ao tamanho da subárvore esquerda (tr[u].l.size), o valor procurado está na subárvore esquerda. Recursivamente, chamamos getKey na subárvore esquerda com o mesmo rank.
  • Se o rank for menor ou igual ao tamanho da subárvore esquerda mais a contagem do nó atual (tr[u].l.size + tr[u].cnt), então o valor procurado é tr[u].key.
  • Caso contrário, o valor procurado está na subárvore direita. Precisamos ajustar o rank subtraindo o tamanho da subárvore esquerda e a contagem do nó atual, e então chamamos getKey recursivamente na subárvore direita.
int getKey(int u, int rank) {
    if (!u) return INF; // Valor sentinela para indicar não encontrado
    if (rank <= tr[tr[u].l].size) return getKey(tr[u].l, rank);
    if (rank <= tr[tr[u].l].size + tr[u].cnt) return tr[u].key;
    return getKey(tr[u].r, rank - tr[tr[u].l].size - tr[u].cnt);
}

Consulta de Ranking por Valor

A função getRank(u, key) retorna o ranking do key na subárvore enraizada em u.

  • Se o key for menor que tr[u].key, o valor procurado está na subárvore esquerda. Chamamos getRank recursivamente na subárvore esquerda.
  • Se o key for igual a tr[u].key, o ranking é o tamanho da subárvore esquerda mais 1 (considerando a contagem do nó atual).
  • Se o key for maior que tr[u].key, o valor procurado está na subárvore direita. O ranking será a soma do tamanho da subárvore esquerda, a contagem do nó atual, e o ranking do key na subárvore direita.
int getRank(int u, int key) {
    if (!u) return 0; // Ranking 0 indica não encontrado
    if (key < tr[u].key) return getRank(tr[u].l, key);
    if (key == tr[u].key) return tr[tr[u].l].size + 1;
    return tr[tr[u].l].size + tr[u].cnt + getRank(tr[u].r, key);
}

Para simplificar o manuseio de casos de borda, frequentemente são adicionados nós sentinela de -INF e +INF. Ao usar sentinelas, a consulta por ranking rank precisa ser ajustada por +1, e a consulta por valor key precisa ter seu resultado subtraído por 1.

Consulta de Antecessor

A função getPrev(u, key) encontra o maior valor na subárvore enraizada em u que é estritamente menor que key.

  • Se tr[u].key for maior ou igual a key, o antecessor, se existir, deve estar na subárvore esquerda. Chamamos getPrev recursivamente na subárvore esquerda.
  • Se tr[u].key for menor que key, o tr[u].key é um candidato a antecessor. Comparamos tr[u].key com o resultado da busca recursiva na subárvore direita e retornamos o maior deles.
int getPrev(int u, int key) {
    if (!u) return -INF; // Valor sentinela
    if (tr[u].key >= key) return getPrev(tr[u].l, key);
    return max(tr[u].key, getPrev(tr[u].r, key));
}

Consulta de Sucesssor

A função getNext(u, key) encontra o menor valor na subárvore enraizada em u que é estritamente maior que key.

  • Se tr[u].key for menor ou igual a key, o sucessor, se existir, deve estar na subárvore direita. Chamamos getNext recursivamente na subárvore direita.
  • Se tr[u].key for maior que key, o tr[u].key é um candidato a sucessor. Comparamos tr[u].key com o resultado da busca recursiva na subárvore esquerda e retornamos o menor deles.
int getNext(int u, int key) {
    if (!u) return INF; // Valor sentinela
    if (tr[u].key <= key) return getNext(tr[u].r, key);
    return min(tr[u].key, getNext(tr[u].l, key));
}

As operações de consulta descritas têm complexidade O(h), onde h é a altura da BST. Em dados aleatórios, h tende a log n. No entanto, BSTs podem degenerar para uma lista ligada (h=n) com entradas ordenadas, tornando a complexidade O(n). Para garantir uma altura logarítmica, são necessárias técnicas de balanceamento.

Rotações

Rotações são operações fundamentais em BSTs balanceadas que preservam a ordem dos elementos (in-order traversal) enquanto modificam a estrutura da árvore para manter o balanceamento. Elas são cruciais para a maioria das árvores balanceadas, exceto por algumas como a FHQ-Treap.

Uma rotação zig (rotação à esquerda) e zag (rotação à direita) reorganizam um nó pai e seu filho para melhorar a altura da árvore.

  • zig(u): Transforma um nó u e seu filho esquerdo q. O filho esquerdo de q se torna o filho direito de u, u se torna o filho direito de q, e q se torna o novo u.
  • zag(u): Transforma um nó u e seu filho direito q. O filho direito de q se torna o filho esquerdo de u, u se torna o filho esquerdo de q, e q se torna o novo u.
void zig(int &u) {
    int q = tr[u].l; // q é o filho esquerdo de u
    tr[u].l = tr[q].r; // O filho direito de q se torna o filho esquerdo de u
    tr[q].r = u;       // u se torna o filho direito de q
    u = q;             // q se torna o novo nó raiz (u)
}

void zag(int &u) {
    int q = tr[u].r; // q é o filho direito de u
    tr[u].r = tr[q].l; // O filho esquerdo de q se torna o filho direito de u
    tr[q].l = u;       // u se torna o filho esquerdo de q
    u = q;             // q se torna o novo nó raiz (u)
}

A passagem por referência (int &u) é importante porque a raiz da subárvore pode mudar após a rotação.

Treap

Treap é uma combinação de Tree (BST) e Heap. A ideia principal é atribuir um valor de prioridade aleatório (val) a cada nó durante a inserção. A árvore mantém a propriedade de BST com base nos key e a propriedade de heap (min-heap ou max-heap) com base nos val. As rotações são usadas para manter a propriedade de heap após a inserção ou deleção.

  • Inserção: Insere o novo nó como em uma BST. Em seguida, usando rotações zig ou zag, o nó é "empurrado" para cima até que a propriedade de heap seja restaurada.
  • Remoção: Para remover um nó, primeiro ele é rotacionado para baixo (como em um heap) até se tornar uma folha, e então é removido.
void insert(int &u, int key) {
    if (!u) u = create(key); // Cria novo nó se a posição estiver vazia
    else if (key == tr[u].key) tr[u].cnt++; // Incrementa contador se a chave já existe
    else if (key < tr[u].key) {
        insert(tr[u].l, key);
        if (tr[tr[u].l].val < tr[u].val) zig(u); // Rotaciona se a propriedade de heap for violada
    } else {
        insert(tr[u].r, key);
        if (tr[tr[u].r].val < tr[u].val) zag(u); // Rotaciona se a propriedade de heap for violada
    }
    pushup(u); // Atualiza o tamanho da subárvore
}

void erase(int &u, int key) {
    if (!u) return; // Nada a fazer se o nó não existir
    if (key == tr[u].key) {
        if (tr[u].cnt > 1) tr[u].cnt--; // Reduz a contagem se houver duplicatas
        else if (tr[u].l || tr[u].r) { // Se não for folha
            // Decide para qual lado rotacionar para descer o nó
            if (!tr[u].r || tr[tr[u].l].val > tr[tr[u].r].val) {
                zig(u);
                erase(tr[u].r, key); // Continua a remoção na subárvore direita
            } else {
                zag(u);
                erase(tr[u].l, key); // Continua a remoção na subárvore esquerda
            }
        } else u = 0; // Remove o nó folha
    } else if (key < tr[u].key) erase(tr[u].l, key); // Busca na subárvore esquerda
    else erase(tr[u].r, key); // Busca na subárvore direita
    
    pushup(u); // Atualiza o tamanho da subárvore
}

A função pushup(u) atualiza o campo size do nó u com base nos tamanhos de suas subárvores esquerda e direita e sua própria contagem cnt.

void pushup(int u) {
    tr[u].size = tr[tr[u].l].size + tr[tr[u].r].size + tr[u].cnt;
}

Com essas operações, a altura da Treap é esperada ser O(log n), resultando em uma complexidade média de O(log n) para todas as operações.

#include <iostream>
#include <cstring>
#include <algorithm>
#include <cstdlib>
#include <ctime> // Para inicializar o gerador de números aleatórios

using namespace std;

const int MAXN = 100005;
const int INF = 0x3f3f3f3f; // Valor sentinela positivo
const int NEG_INF = -0x3f3f3f3f; // Valor sentinela negativo

struct Node {
    int l, r;       // Ponteiros para filhos esquerdo e direito
    int key;        // Valor do nó
    int val;        // Prioridade aleatória (para heap)
    int cnt;        // Contagem de ocorrências da chave
    int size;       // Tamanho da subárvore
} tr[MAXN];

int root; // Raiz da Treap
int node_count; // Contador para alocação de nós

// Função para criar um novo nó
int createNode(int key) {
    node_count++;
    tr[node_count].key = key;
    tr[node_count].val = rand(); // Atribui prioridade aleatória
    tr[node_count].cnt = 1;
    tr[node_count].size = 1;
    tr[node_count].l = tr[node_count].r = 0; // Inicializa filhos como nulos
    return node_count;
}

// Atualiza o tamanho da subárvore de um nó
void pushUp(int u) {
    tr[u].size = tr[tr[u].l].size + tr[tr[u].r].size + tr[u].cnt;
}

// Rotação à esquerda (zig)
void zig(int &u) {
    int q = tr[u].l;
    tr[u].l = tr[q].r;
    tr[q].r = u;
    pushUp(u); // Atualiza u primeiro
    u = q;     // Atualiza a raiz da subárvore
    pushUp(u); // Atualiza a nova raiz
}

// Rotação à direita (zag)
void zag(int &u) {
    int q = tr[u].r;
    tr[u].r = tr[q].l;
    tr[q].l = u;
    pushUp(u); // Atualiza u primeiro
    u = q;     // Atualiza a raiz da subárvore
    pushUp(u); // Atualiza a nova raiz
}

// Inicializa a árvore com nós sentinela
void build() {
    node_count = 0; // Reseta contador de nós
    // Cria nós sentinela -INF e +INF
    int negInfNode = createNode(NEG_INF);
    int posInfNode = createNode(INF);
    root = negInfNode;
    tr[root].r = posInfNode; // -INF aponta para +INF como filho direito
    pushUp(root); // Atualiza o tamanho do nó -INF
}

// Insere uma chave na árvore
void insert(int &u, int key) {
    if (!u) {
        u = createNode(key);
        return;
    }
    if (key == tr[u].key) {
        tr[u].cnt++; // Incrementa contador se a chave já existe
    } else if (key < tr[u].key) {
        insert(tr[u].l, key);
        // Se a prioridade do filho esquerdo for menor, rotaciona para cima
        if (tr[tr[u].l].val < tr[u].val) zig(u);
    } else {
        insert(tr[u].r, key);
        // Se a prioridade do filho direito for menor, rotaciona para cima
        if (tr[tr[u].r].val < tr[u].val) zag(u);
    }
    pushUp(u); // Atualiza o tamanho após a inserção/rotação
}

// Remove uma chave da árvore
void erase(int &u, int key) {
    if (!u) return; // Nó não encontrado
    
    if (key == tr[u].key) {
        if (tr[u].cnt > 1) {
            tr[u].cnt--; // Reduz a contagem se houver duplicatas
        } else if (tr[u].l || tr[u].r) { // Se não for folha
            // Escolhe para qual lado rotacionar para descer o nó, priorizando o filho com maior prioridade
            if (!tr[u].r || (tr[u].l && tr[tr[u].l].val > tr[tr[u].r].val)) {
                zig(u);
                erase(tr[u].r, key); // Continua a remoção na subárvore direita
            } else {
                zag(u);
                erase(tr[u].l, key); // Continua a remoção na subárvore esquerda
            }
        } else {
            u = 0; // Remove o nó folha
        }
    } else if (key < tr[u].key) {
        erase(tr[u].l, key); // Busca na subárvore esquerda
    } else {
        erase(tr[u].r, key); // Busca na subárvore direita
    }
    
    // Se o nó não foi completamente removido (u != 0), atualiza seu tamanho
    if (u != 0) pushUp(u);
}

// Consulta o ranking de uma chave
int getRank(int u, int key) {
    if (!u) return 0; // Chave não encontrada
    if (key < tr[u].key) {
        return getRank(tr[u].l, key); // Busca na subárvore esquerda
    } else if (key == tr[u].key) {
        return tr[tr[u].l].size + 1; // Ranking é tamanho da subárvore esquerda + 1
    } else {
        // Ranking é tamanho da subárvore esquerda + contagem do nó atual + ranking na subárvore direita
        return tr[tr[u].l].size + tr[u].cnt + getRank(tr[u].r, key);
    }
}

// Consulta a chave com um dado ranking
int getKey(int u, int rank) {
    if (!u) return INF; // Ranking inválido (fora dos limites)
    if (rank <= tr[tr[u].l].size) {
        return getKey(tr[u].l, rank); // Busca na subárvore esquerda
    } else if (rank <= tr[tr[u].l].size + tr[u].cnt) {
        return tr[u].key; // A chave está neste nó
    } else {
        // Busca na subárvore direita com rank ajustado
        return getKey(tr[u].r, rank - tr[tr[u].l].size - tr[u].cnt);
    }
}

// Consulta o antecessor de uma chave
int getPrev(int u, int key) {
    if (!u) return NEG_INF; // Antecessor não encontrado
    if (tr[u].key >= key) {
        return getPrev(tr[u].l, key); // Antecessor está na subárvore esquerda
    } else {
        // Compara o valor atual com o antecessor na subárvore direita
        return max(tr[u].key, getPrev(tr[u].r, key));
    }
}

// Consulta o sucessor de uma chave
int getNext(int u, int key) {
    if (!u) return INF; // Sucessor não encontrado
    if (tr[u].key <= key) {
        return getNext(tr[u].r, key); // Sucessor está na subárvore direita
    } else {
        // Compara o valor atual com o sucessor na subárvore esquerda
        return min(tr[u].key, getNext(tr[u].l, key));
    }
}

int main() {
    srand(time(0)); // Inicializa o gerador de números aleatórios

    int n;
    scanf("%d", &n);
    
    build(); // Inicializa a Treap com sentinelas

    while (n--) {
        int op, x;
        scanf("%d%d", &op, &x);
        switch (op) {
            case 1: // Inserir
                insert(root, x);
                break;
            case 2: // Remover
                erase(root, x);
                break;
            case 3: // Consultar ranking por valor
                // Subtrai 1 do resultado porque o ranking base (com sentinela) é 1-based
                printf("%d\n", getRank(root, x) - 1);
                break;
            case 4: // Consultar valor por ranking
                // Adiciona 1 ao ranking para buscar o elemento correto após as sentinelas
                printf("%d\n", getKey(root, x + 1));
                break;
            case 5: // Consultar antecessor
                printf("%d\n", getPrev(root, x));
                break;
            case 6: // Consultar sucessor
                printf("%d\n", getNext(root, x));
                break;
        }
    }
    
    return 0;
}

Tags: Treap Árvore de Busca Binária balanceamento de árvore estrutura de dados heap

Publicado em 7-18 19:23