Trie (Árvore de Prefixos) em Zig — Implementação Completa

Trie (Árvore de Prefixos) em Zig — Implementação Completa

A Trie é uma estrutura de dados em forma de árvore otimizada para armazenar e buscar strings. Cada nó representa um caractere, e caminhos da raiz formam palavras. Busca, inserção e remoção operam em O(m), onde m é o comprimento da string.

Conceito

Palavras: "zig", "zip", "zona", "zero", "ziglang"

             (raiz)
            /       \
           z          ...
          / \
         i    o
        / \    \
       g*  p*   n
       |        |
       l        a*
       |
       a
       |
       n
       |
       g*

* = fim de palavra

Memória: cada nó tem array de 26 ponteiros (letras minúsculas)
Compressão: Trie compacta (Patricia/Radix) agrupa prefixos comuns

Implementação em Zig

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

const ALPHABET_SIZE = 26;

pub const Trie = struct {
    pub const No = struct {
        filhos: [ALPHABET_SIZE]?*No,
        fim_palavra: bool,
        contagem_prefixo: u32,
        contagem_palavra: u32,
    };

    raiz: *No,
    total_palavras: u32,
    allocator: Allocator,

    pub fn init(allocator: Allocator) !Trie {
        const raiz = try criarNo(allocator);
        return .{ .raiz = raiz, .total_palavras = 0, .allocator = allocator };
    }

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

    fn criarNo(allocator: Allocator) !*No {
        const no = try allocator.create(No);
        no.* = .{
            .filhos = [_]?*No{null} ** ALPHABET_SIZE,
            .fim_palavra = false,
            .contagem_prefixo = 0,
            .contagem_palavra = 0,
        };
        return no;
    }

    fn liberarNo(allocator: Allocator, no: *No) void {
        for (no.filhos) |filho| {
            if (filho) |f| liberarNo(allocator, f);
        }
        allocator.destroy(no);
    }

    /// Insere uma palavra.
    pub fn inserir(self: *Trie, palavra: []const u8) !void {
        var atual = self.raiz;
        for (palavra) |c| {
            const idx = c - 'a';
            if (atual.filhos[idx] == null) {
                atual.filhos[idx] = try criarNo(self.allocator);
            }
            atual = atual.filhos[idx].?;
            atual.contagem_prefixo += 1;
        }
        atual.fim_palavra = true;
        atual.contagem_palavra += 1;
        self.total_palavras += 1;
    }

    /// Busca exata.
    pub fn buscar(self: *const Trie, palavra: []const u8) bool {
        const no = self.navegarAte(palavra);
        return if (no) |n| n.fim_palavra else false;
    }

    /// Verifica se existe alguma palavra com este prefixo.
    pub fn temPrefixo(self: *const Trie, prefixo: []const u8) bool {
        return self.navegarAte(prefixo) != null;
    }

    /// Conta palavras com este prefixo.
    pub fn contarPrefixo(self: *const Trie, prefixo: []const u8) u32 {
        const no = self.navegarAte(prefixo);
        return if (no) |n| n.contagem_prefixo else 0;
    }

    /// Remove uma palavra (decrementa contadores).
    pub fn remover(self: *Trie, palavra: []const u8) bool {
        if (!self.buscar(palavra)) return false;

        var atual = self.raiz;
        for (palavra) |c| {
            atual = atual.filhos[c - 'a'].?;
            atual.contagem_prefixo -= 1;
        }
        atual.contagem_palavra -= 1;
        if (atual.contagem_palavra == 0) atual.fim_palavra = false;
        self.total_palavras -= 1;
        return true;
    }

    /// Autocompletar: lista palavras com o prefixo dado.
    pub fn autocompletar(
        self: *const Trie,
        allocator: Allocator,
        prefixo: []const u8,
        max: usize,
    ) ![][]u8 {
        var resultados = std.ArrayList([]u8).init(allocator);
        const no = self.navegarAte(prefixo);
        if (no == null) return try resultados.toOwnedSlice();

        var buffer = std.ArrayList(u8).init(allocator);
        defer buffer.deinit();
        try buffer.appendSlice(prefixo);

        try coletarPalavras(no.?, &buffer, &resultados, allocator, max);
        return try resultados.toOwnedSlice();
    }

    fn navegarAte(self: *const Trie, chave: []const u8) ?*No {
        var atual = self.raiz;
        for (chave) |c| {
            const idx = c - 'a';
            if (atual.filhos[idx]) |filho| {
                atual = filho;
            } else return null;
        }
        return atual;
    }

    fn coletarPalavras(no: *No, buf: *std.ArrayList(u8), res: *std.ArrayList([]u8), alloc: Allocator, max: usize) !void {
        if (res.items.len >= max) return;
        if (no.fim_palavra) {
            try res.append(try alloc.dupe(u8, buf.items));
        }
        for (0..ALPHABET_SIZE) |i| {
            if (no.filhos[i]) |filho| {
                try buf.append(@intCast('a' + i));
                try coletarPalavras(filho, buf, res, alloc, max);
                _ = buf.pop();
            }
        }
    }
};

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

    var trie = try Trie.init(allocator);
    defer trie.deinit();

    const palavras = [_][]const u8{ "zig", "zip", "zona", "zero", "ziglang", "zigzag" };
    for (palavras) |p| try trie.inserir(p);

    try stdout.print("Total: {d}\n", .{trie.total_palavras});
    try stdout.print("Buscar 'zig': {}\n", .{trie.buscar("zig")});
    try stdout.print("Prefixo 'zi': {d} palavras\n", .{trie.contarPrefixo("zi")});

    const sug = try trie.autocompletar(allocator, "zi", 10);
    defer {
        for (sug) |s| allocator.free(s);
        allocator.free(sug);
    }
    try stdout.print("Autocompletar 'zi':\n", .{});
    for (sug) |s| try stdout.print("  {s}\n", .{s});
}

Análise de Complexidade

OperaçãoTempoEspaço
InserirO(m)O(m * ALPHABET)
BuscarO(m)O(1)
PrefixoO(m)O(1)
RemoverO(m)O(1)

Recursos Relacionados

Continue aprendendo Zig

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