Árvore Red-Black em Zig — Implementação Completa

Árvore Red-Black em Zig — Implementação Completa

A Red-Black Tree (árvore rubro-negra) é uma BST auto-balanceada onde cada nó tem uma cor (vermelho ou preto) e regras específicas garantem que a árvore nunca fica muito desbalanceada.

Conceito

Regras

  1. Todo nó é vermelho ou preto
  2. A raiz é sempre preta
  3. Nós nulos (folhas NIL) são pretos
  4. Se um nó é vermelho, seus filhos são pretos
  5. Todo caminho de um nó até suas folhas NIL tem o mesmo número de nós pretos
        8(B)
       /    \
     4(R)   12(R)
    /   \   /   \
  2(B) 6(B) 10(B) 14(B)
  / \  / \  / \   / \
 1  3 5  7 9 11 13  15

B = preto, R = vermelho
Caminho raiz→folha: sempre 2 nós pretos (black-height = 2)

Inserção e correção de cores:
  Inserir sempre como VERMELHO → verificar violações → rotações e recolorir

Implementação em Zig

const std = @import("std");
const Allocator = std.mem.Allocator;

pub fn RedBlackTree(comptime T: type) type {
    return struct {
        const Self = @This();
        const Cor = enum { vermelho, preto };

        pub const No = struct {
            valor: T,
            esquerdo: ?*No,
            direito: ?*No,
            pai: ?*No,
            cor: Cor,
        };

        raiz: ?*No,
        tamanho: usize,
        allocator: Allocator,

        pub fn init(allocator: Allocator) Self {
            return .{ .raiz = null, .tamanho = 0, .allocator = allocator };
        }

        pub fn deinit(self: *Self) void {
            liberarNo(self.allocator, self.raiz);
        }

        fn liberarNo(allocator: Allocator, no: ?*No) void {
            if (no) |n| {
                liberarNo(allocator, n.esquerdo);
                liberarNo(allocator, n.direito);
                allocator.destroy(n);
            }
        }

        fn corDe(no: ?*const No) Cor {
            return if (no) |n| n.cor else .preto;
        }

        fn rotacaoEsquerda(self: *Self, x: *No) void {
            const y = x.direito.?;
            x.direito = y.esquerdo;
            if (y.esquerdo) |yl| yl.pai = x;
            y.pai = x.pai;
            if (x.pai == null) {
                self.raiz = y;
            } else if (x.pai.?.esquerdo == x) {
                x.pai.?.esquerdo = y;
            } else {
                x.pai.?.direito = y;
            }
            y.esquerdo = x;
            x.pai = y;
        }

        fn rotacaoDireita(self: *Self, y: *No) void {
            const x = y.esquerdo.?;
            y.esquerdo = x.direito;
            if (x.direito) |xr| xr.pai = y;
            x.pai = y.pai;
            if (y.pai == null) {
                self.raiz = x;
            } else if (y.pai.?.esquerdo == y) {
                y.pai.?.esquerdo = x;
            } else {
                y.pai.?.direito = x;
            }
            x.direito = y;
            y.pai = x;
        }

        /// Insere um valor — O(log n).
        pub fn inserir(self: *Self, valor: T) !void {
            const novo = try self.allocator.create(No);
            novo.* = .{
                .valor = valor,
                .esquerdo = null,
                .direito = null,
                .pai = null,
                .cor = .vermelho,
            };

            // BST insert
            var pai: ?*No = null;
            var atual = self.raiz;
            while (atual) |a| {
                pai = a;
                if (valor < a.valor) {
                    atual = a.esquerdo;
                } else if (valor > a.valor) {
                    atual = a.direito;
                } else {
                    self.allocator.destroy(novo);
                    return; // duplicata
                }
            }

            novo.pai = pai;
            if (pai == null) {
                self.raiz = novo;
            } else if (valor < pai.?.valor) {
                pai.?.esquerdo = novo;
            } else {
                pai.?.direito = novo;
            }

            self.tamanho += 1;
            self.corrigirInsercao(novo);
        }

        fn corrigirInsercao(self: *Self, z_param: *No) void {
            var z = z_param;
            while (z.pai != null and z.pai.?.cor == .vermelho) {
                if (z.pai == z.pai.?.pai.?.esquerdo) {
                    const tio = z.pai.?.pai.?.direito;
                    if (corDe(tio) == .vermelho) {
                        // Caso 1: tio vermelho
                        z.pai.?.cor = .preto;
                        if (tio) |t| t.cor = .preto;
                        z.pai.?.pai.?.cor = .vermelho;
                        z = z.pai.?.pai.?;
                    } else {
                        if (z == z.pai.?.direito) {
                            // Caso 2: triângulo
                            z = z.pai.?;
                            self.rotacaoEsquerda(z);
                        }
                        // Caso 3: linha
                        z.pai.?.cor = .preto;
                        z.pai.?.pai.?.cor = .vermelho;
                        self.rotacaoDireita(z.pai.?.pai.?);
                    }
                } else {
                    const tio = z.pai.?.pai.?.esquerdo;
                    if (corDe(tio) == .vermelho) {
                        z.pai.?.cor = .preto;
                        if (tio) |t| t.cor = .preto;
                        z.pai.?.pai.?.cor = .vermelho;
                        z = z.pai.?.pai.?;
                    } else {
                        if (z == z.pai.?.esquerdo) {
                            z = z.pai.?;
                            self.rotacaoDireita(z);
                        }
                        z.pai.?.cor = .preto;
                        z.pai.?.pai.?.cor = .vermelho;
                        self.rotacaoEsquerda(z.pai.?.pai.?);
                    }
                }
            }
            self.raiz.?.cor = .preto;
        }

        /// Busca — O(log n).
        pub fn buscar(self: *const Self, valor: T) bool {
            var atual = self.raiz;
            while (atual) |no| {
                if (valor == no.valor) return true;
                atual = if (valor < no.valor) no.esquerdo else no.direito;
            }
            return false;
        }

        /// Travessia in-order.
        pub fn inOrder(self: *const Self, allocator: Allocator) ![]T {
            var res = std.ArrayList(T).init(allocator);
            inOrderHelper(self.raiz, &res) catch {};
            return try res.toOwnedSlice();
        }

        fn inOrderHelper(no: ?*const No, res: *std.ArrayList(T)) !void {
            if (no == null) return;
            const n = no.?;
            try inOrderHelper(n.esquerdo, res);
            try res.append(n.valor);
            try inOrderHelper(n.direito, res);
        }

        /// Black-height (número de nós pretos até folha).
        pub fn blackHeight(self: *const Self) u32 {
            var h: u32 = 0;
            var atual = self.raiz;
            while (atual) |no| {
                if (no.cor == .preto) h += 1;
                atual = no.esquerdo;
            }
            return h;
        }
    };
}

pub fn main() !void {
    const stdout = std.io.getStdOut().writer();
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    var rb = RedBlackTree(i32).init(allocator);
    defer rb.deinit();

    for ([_]i32{ 7, 3, 18, 10, 22, 8, 11, 26, 2, 6, 13 }) |v| {
        try rb.inserir(v);
    }

    try stdout.print("Tamanho: {d}\n", .{rb.tamanho});
    try stdout.print("Raiz: {d} (cor: {s})\n", .{
        rb.raiz.?.valor,
        if (rb.raiz.?.cor == .preto) "preto" else "vermelho",
    });
    try stdout.print("Black-height: {d}\n", .{rb.blackHeight()});

    const ordenado = try rb.inOrder(allocator);
    defer allocator.free(ordenado);
    try stdout.print("In-order: ", .{});
    for (ordenado) |v| try stdout.print("{d} ", .{v});
    try stdout.print("\n", .{});

    try stdout.print("Buscar 10: {}\n", .{rb.buscar(10)});
    try stdout.print("Buscar 15: {}\n", .{rb.buscar(15)});
}

Análise de Complexidade

OperaçãoTempoEspaço
BuscaO(log n)O(1)
InserçãoO(log n)O(1)
RemoçãoO(log n)O(1)

Recursos Relacionados

Continue aprendendo Zig

Explore mais tutoriais e artigos em português para dominar a linguagem Zig.