Implementación de Árbol Rojo-Negro en C++

Propiedades del Árbol Rojo-Negro

Un árbol rojo-negro es una estructura de datos que satisface las siguientes invariantes:

  1. Es un árbol binario de búsqueda válido.
  2. Cada nodo tiene asignado un color: rojo o negro.
  3. La raíz y los nodos nulos (hojas NIL) son siempre negros.
  4. Los hijos de un nodo rojo deben ser obligatoriamente negros.
  5. Todas las rutas desde la raíz hasta cualquier hoja NIL contienen la misma cantidad de nodos negros (altura negra).

Estas restricciones garantizan que la altura del árbol está acotada por el doble de la altura negra, manteniendo así una complejidad de \\(O(\log n)\\).

Rotaciones

Las rotaciones son operaciones estructurales que preservan el orden del árbol binario de búsqueda. Una rotación derecha mueve el hijo izquierdo hacia arriba, mientras que una rotación izquierda hace lo inverso.

void rotar(int nodo) {
    int padre = arbol[nodo].padre;
    int abuelo = arbol[padre].padre;
    int dir = (arbol[padre].hijo[1] == nodo) ? 1 : 0;
    int sub = arbol[nodo].hijo[dir ^ 1];
    
    if (sub != 0) arbol[sub].padre = padre;
    arbol[padre].hijo[dir] = sub;
    arbol[nodo].hijo[dir ^ 1] = padre;
    
    if (abuelo != 0)
        arbol[abuelo].hijo[(arbol[abuelo].hijo[1] == padre) ? 1 : 0] = nodo;
    else
        raiz = nodo;
    
    arbol[padre].padre = nodo;
    arbol[nodo].padre = abuelo;
    actualizar(padre);
    actualizar(nodo);
}

Inserción

Si el árbol está vacío, se crea un nodo negro como raíz. En caso contrario, el nuevo nodo se inserta de color rojo, lo que únicamente puede violar la propiedad de que dos nodos rojos no pueden ser adyacentes. Para restaurar el equilibrio se aplica un procedimiento llamado corrección de doble rojo.

int insertar(int valor) {
    if (raiz == 0) {
        raiz = crearNodo(valor, NEGRO);
        return raiz;
    }
    int actual = raiz;
    while (true) {
        arbol[actual].tam++;
        if (arbol[actual].valor == valor) {
            arbol[actual].repeticion++;
            return actual;
        }
        int lado = (arbol[actual].valor < valor) ? 1 : 0;
        if (arbol[actual].hijo[lado] == 0) {
            int nuevo = crearNodo(valor, ROJO);
            arbol[nuevo].padre = actual;
            arbol[actual].hijo[lado] = nuevo;
            corregirDobleRojo(nuevo);
            return nuevo;
        }
        actual = arbol[actual].hijo[lado];
    }
}

Corrección de Doble Rojo

La función corregirDobleRojo(nodo) asume que el subárbol bajo nodo ya cumple las propiedades, pero nodo y su padre podrían ser ambos rojos.

Sea n el nodo a corregir, p su padre, g su abuelo, y t el tío de n.

  • Si t es negro (o NIL): Se realizan rotaciones para colocar el nodo intermedio como raíz del subárbol. Si n y p están en la misma dirección, se rota p y se intercambian colores de p y g. Si están en direcciones opuestas, primero se rota n, equiparando al caso anterior.
  • Si t es rojo: Se colorean p y t de negro, g de rojo, y se recursa sobre g.
void corregirDobleRojo(int nodo) {
    if (nodo == raiz || arbol[nodo].padre == raiz) {
        arbol[raiz].color = NEGRO;
        return;
    }
    int p = arbol[nodo].padre;
    if (arbol[p].color == NEGRO) return;
    
    int g = arbol[p].padre;
    int t = arbol[g].hijo[(arbol[g].hijo[0] == p) ? 1 : 0];
    
    if (arbol[t].color == NEGRO) {
        bool dirN = (arbol[p].hijo[1] == nodo);
        bool dirP = (arbol[g].hijo[1] == p);
        if (dirN != dirP) {
            rotar(nodo);
            int tmp = nodo; nodo = p; p = tmp;
        }
        rotar(p);
        arbol[g].color = ROJO;
        arbol[p].color = NEGRO;
        return;
    }
    
    arbol[p].color = arbol[t].color = NEGRO;
    arbol[g].color = ROJO;
    corregirDobleRojo(g);
}

Eliminación

Para eliminar un valor, se localiza el nodo correspondiente. Si el nodo tiene dos hijos, se intercambia con su sucesor in-order, de modo que la eliminación siempre ocurre en un nodo con a lo sumo un hijo. Si el nodo a eliminar es negro, se requiere una corrección de doble negro antes de removerlo.

void eliminar(int valor) {
    int nodo = raiz;
    while (arbol[nodo].valor != valor && arbol[nodo].hijo[arbol[nodo].valor < valor])
        nodo = arbol[nodo].hijo[arbol[nodo].valor < valor];
    
    // Actualizar tamaños en el camino
    for (int i = raiz; i != 0 && arbol[i].valor != valor; 
         i = arbol[i].hijo[arbol[i].valor < valor])
        arbol[i].tam--;
    
    if (arbol[nodo].valor != valor) return;
    
    arbol[nodo].tam--;
    if (--arbol[nodo].repeticion > 0) return;
    
    // Caso: árbol queda vacío
    if (nodo == raiz && !arbol[nodo].hijo[0] && !arbol[nodo].hijo[1]) {
        raiz = 0;
        reciclar(nodo);
        return;
    }
    
    // Buscar sucesor y mover datos
    int objetivo = nodo;
    while (arbol[objetivo].hijo[0] || arbol[objetivo].hijo[1]) {
        int suc;
        if (!arbol[objetivo].hijo[0]) suc = arbol[objetivo].hijo[1];
        else if (!arbol[objetivo].hijo[1]) suc = arbol[objetivo].hijo[0];
        else {
            suc = arbol[objetivo].hijo[1];
            while (arbol[suc].hijo[0]) suc = arbol[suc].hijo[0];
        }
        swap(arbol[objetivo].valor, arbol[suc].valor);
        swap(arbol[objetivo].repeticion, arbol[suc].repeticion);
        objetivo = suc;
    }
    
    if (arbol[objetivo].color == NEGRO)
        corregirDobleNegro(objetivo);
    
    int padreObj = arbol[objetivo].padre;
    if (padreObj)
        arbol[padreObj].hijo[(arbol[padreObj].hijo[1] == objetivo) ? 1 : 0] = 0;
    
    reciclar(objetivo);
    for (int i = padreObj; i; i = arbol[i].padre)
        actualizar(i);
}

Corrección de Doble Negro

Cuando se elimina un nodo negro, la altura negra de esa rama disminuye. La función corregirDobleNegro(nodo) aborda este desbalance. Sea s el hermano del nodo problemático y p su padre.

  • Si s es rojo: Se rota s y se intercambian colores de s y p, reduciendo al caso donde el hermano es negro.
  • Si s es negro con un hijo rojo w: Se realizan rotaciones posicionando el nodo intermedio como raíz. Se conserva el color original de p y se pintan s y w de negro.
  • Si s es negro sin hijos rojos: Se pinta s de rojo. Si p es rojo, se pinta de negro y termina. Si p es negro, se recursa sobre p.
void corregirDobleNegro(int nodo) {
    if (nodo == raiz) {
        arbol[nodo].color = NEGRO;
        return;
    }
    
    int p = arbol[nodo].padre;
    int ladoN = (arbol[p].hijo[0] == nodo) ? 0 : 1;
    int s = arbol[p].hijo[ladoN ^ 1];
    
    if (arbol[s].color == ROJO) {
        arbol[s].color = NEGRO;
        arbol[p].color = ROJO;
        rotar(s);
        p = arbol[nodo].padre;
        ladoN = (arbol[p].hijo[0] == nodo) ? 0 : 1;
        s = arbol[p].hijo[ladoN ^ 1];
    }
    
    int cercano = arbol[s].hijo[ladoN];
    int lejano = arbol[s].hijo[ladoN ^ 1];
    
    if (arbol[cercano].color == ROJO) {
        rotar(cercano);
        s = cercano;
        cercano = arbol[s].hijo[ladoN];
        lejano = arbol[s].hijo[ladoN ^ 1];
    }
    
    if (arbol[lejano].color == ROJO) {
        arbol[s].color = arbol[p].color;
        arbol[p].color = NEGRO;
        arbol[lejano].color = NEGRO;
        rotar(s);
        return;
    }
    
    // Ningún hijo rojo
    arbol[s].color = ROJO;
    if (arbol[p].color == ROJO) {
        arbol[p].color = NEGRO;
        return;
    }
    corregirDobleNegro(p);
}

Operaciones de Consulta

Las consultas estándar —rengo, k-ésimo elemento, predecesor y sucesor— se implementan de forma idéntica a un árbol binario de búsqueda convencional, aprovechando el campo de tamaño almacenado en cada nodo.

int consultaRango(int valor) {
    int resultado = 1, actual = raiz;
    while (actual) {
        if (arbol[actual].valor == valor)
            return resultado + arbol[arbol[actual].hijo[0]].tam;
        if (valor < arbol[actual].valor)
            actual = arbol[actual].hijo[0];
        else {
            resultado += arbol[arbol[actual].hijo[0]].tam + arbol[actual].repeticion;
            actual = arbol[actual].hijo[1];
        }
    }
    return resultado;
}

int consultaKesimo(int k) {
    int acumulado = 0, actual = raiz;
    while (true) {
        int tamIzq = arbol[arbol[actual].hijo[0]].tam;
        if (acumulado + tamIzq < k && k <= acumulado + tamIzq + arbol[actual].repeticion)
            return arbol[actual].valor;
        if (k <= acumulado + tamIzq)
            actual = arbol[actual].hijo[0];
        else {
            acumulado += tamIzq + arbol[actual].repeticion;
            actual = arbol[actual].hijo[1];
        }
    }
}

int predecesor(int valor) {
    int mejor = 0, actual = raiz;
    while (actual) {
        if (arbol[actual].valor < valor) {
            mejor = arbol[actual].valor;
            actual = arbol[actual].hijo[1];
        } else
            actual = arbol[actual].hijo[0];
    }
    return mejor;
}

int sucesor(int valor) {
    int mejor = 0, actual = raiz;
    while (actual) {
        if (arbol[actual].valor > valor) {
            mejor = arbol[actual].valor;
            actual = arbol[actual].hijo[0];
        } else
            actual = arbol[actual].hijo[1];
    }
    return mejor;
}

Código Completo

#include <bits>
using namespace std;

const int MAXN = 1100005;
const int ROJO = 1;
const int NEGRO = 0;

struct Nodo {
    int padre;
    int hijo[2];
    int tam;
    int repeticion;
    int valor;
    int color;
};

Nodo arbol[MAXN];
queue<int> poolReciclaje;
int contadorNodos = 0;
int raiz = 0;

int leer() {
    int x = 0, c = getchar(), signo = 1;
    while (c < '0' || c > '9') { if (c == '-') signo = -1; c = getchar(); }
    while (c >= '0' && c <= '9') { x = x * 10 + (c - '0'); c = getchar(); }
    return x * signo;
}

int crearNodo(int v, int col) {
    int idx;
    if (poolReciclaje.empty()) idx = ++contadorNodos;
    else { idx = poolReciclaje.front(); poolReciclaje.pop(); }
    arbol[idx] = {0, 0, 0, 1, 1, v, col};
    return idx;
}

void reciclar(int idx) { poolReciclaje.push(idx); }

void actualizar(int nodo) {
    arbol[nodo].tam = arbol[arbol[nodo].hijo[0]].tam 
                    + arbol[arbol[nodo].hijo[1]].tam 
                    + arbol[nodo].repeticion;
}

void rotar(int nodo) {
    int p = arbol[nodo].padre;
    int g = arbol[p].padre;
    int dir = (arbol[p].hijo[1] == nodo) ? 1 : 0;
    int sub = arbol[nodo].hijo[dir ^ 1];
    if (sub) arbol[sub].padre = p;
    arbol[p].hijo[dir] = sub;
    arbol[nodo].hijo[dir ^ 1] = p;
    if (g) arbol[g].hijo[(arbol[g].hijo[1] == p) ? 1 : 0] = nodo;
    else raiz = nodo;
    arbol[p].padre = nodo;
    arbol[nodo].padre = g;
    actualizar(p);
    actualizar(nodo);
}

void corregirDobleRojo(int nodo) {
    if (nodo == raiz || arbol[nodo].padre == raiz) {
        arbol[raiz].color = NEGRO;
        return;
    }
    int p = arbol[nodo].padre;
    if (arbol[p].color == NEGRO) return;
    int g = arbol[p].padre;
    int t = arbol[g].hijo[(arbol[g].hijo[0] == p) ? 1 : 0];
    if (arbol[t].color == NEGRO) {
        bool dn = (arbol[p].hijo[1] == nodo);
        bool dp = (arbol[g].hijo[1] == p);
        if (dn != dp) { rotar(nodo); swap(nodo, p); }
        rotar(p);
        arbol[g].color = ROJO;
        arbol[p].color = NEGRO;
        return;
    }
    arbol[p].color = arbol[t].color = NEGRO;
    arbol[g].color = ROJO;
    corregirDobleRojo(g);
}

void corregirDobleNegro(int nodo) {
    if (nodo == raiz) { arbol[nodo].color = NEGRO; return; }
    int p = arbol[nodo].padre;
    int lado = (arbol[p].hijo[0] == nodo) ? 0 : 1;
    int s = arbol[p].hijo[lado ^ 1];
    if (arbol[s].color == ROJO) {
        arbol[s].color = NEGRO;
        arbol[p].color = ROJO;
        rotar(s);
        p = arbol[nodo].padre;
        lado = (arbol[p].hijo[0] == nodo) ? 0 : 1;
        s = arbol[p].hijo[lado ^ 1];
    }
    int cercano = arbol[s].hijo[lado];
    int lejano = arbol[s].hijo[lado ^ 1];
    if (arbol[cercano].color == ROJO) {
        rotar(cercano);
        s = cercano;
        lejano = arbol[s].hijo[lado ^ 1];
    }
    if (arbol[lejano].color == ROJO) {
        arbol[s].color = arbol[p].color;
        arbol[p].color = NEGRO;
        arbol[lejano].color = NEGRO;
        rotar(s);
        return;
    }
    arbol[s].color = ROJO;
    if (arbol[p].color == ROJO) { arbol[p].color = NEGRO; return; }
    corregirDobleNegro(p);
}

void insertar(int valor) {
    if (!raiz) { raiz = crearNodo(valor, NEGRO); return; }
    int actual = raiz;
    while (true) {
        arbol[actual].tam++;
        if (arbol[actual].valor == valor) {
            arbol[actual].repeticion++;
            return;
        }
        int lado = (arbol[actual].valor < valor) ? 1 : 0;
        if (!arbol[actual].hijo[lado]) {
            int nuevo = crearNodo(valor, ROJO);
            arbol[nuevo].padre = actual;
            arbol[actual].hijo[lado] = nuevo;
            corregirDobleRojo(nuevo);
            return;
        }
        actual = arbol[actual].hijo[lado];
    }
}

void eliminar(int valor) {
    int nodo = raiz;
    while (arbol[nodo].valor != valor && arbol[nodo].hijo[arbol[nodo].valor < valor])
        nodo = arbol[nodo].hijo[arbol[nodo].valor < valor];
    for (int i = raiz; i && arbol[i].valor != valor; i = arbol[i].hijo[arbol[i].valor < valor])
        arbol[i].tam--;
    if (arbol[nodo].valor != valor) return;
    arbol[nodo].tam--;
    if (--arbol[nodo].repeticion) return;
    if (nodo == raiz && !arbol[nodo].hijo[0] && !arbol[nodo].hijo[1]) {
        raiz = 0; reciclar(nodo); return;
    }
    int obj = nodo;
    while (arbol[obj].hijo[0] || arbol[obj].hijo[1]) {
        int suc;
        if (!arbol[obj].hijo[0]) suc = arbol[obj].hijo[1];
        else if (!arbol[obj].hijo[1]) suc = arbol[obj].hijo[0];
        else { suc = arbol[obj].hijo[1]; while (arbol[suc].hijo[0]) suc = arbol[suc].hijo[0]; }
        swap(arbol[obj].valor, arbol[suc].valor);
        swap(arbol[obj].repeticion, arbol[suc].repeticion);
        obj = suc;
    }
    reciclar(obj);
    if (arbol[obj].color == NEGRO) corregirDobleNegro(obj);
    int p = arbol[obj].padre;
    if (p) arbol[p].hijo[(arbol[p].hijo[1] == obj) ? 1 : 0] = 0;
    for (int i = p; i; i = arbol[i].padre) actualizar(i);
}

int main() {
    int n = leer(), m = leer();
    for (int i = 0; i < n; i++) insertar(leer());
    int xorAcum = 0, ult = 0;
    for (int i = 0; i < m; i++) {
        int op = leer(), x = leer() ^ ult;
        if (op == 1) insertar(x);
        else if (op == 2) eliminar(x);
        else if (op == 3) { ult = consultaRango(x); xorAcum ^= ult; }
        else if (op == 4) { ult = consultaKesimo(x); xorAcum ^= ult; }
        else if (op == 5) { ult = predecesor(x); xorAcum ^= ult; }
        else { ult = sucesor(x); xorAcum ^= ult; }
    }
    printf("%d\n", xorAcum);
    return 0;
}</int></bits>

Etiquetas: Red-Black Tree Árbol Balanceado Estructura de Datos C++ Árbol Binario de Búsqueda

Publicado el 7-20 08:57