Inserindo: "zig", "zip", "zona", "zero"

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

   z - e - r - o*

Legenda: * = fim de palavra

Buscar "zig":  z → i → g* → encontrado!
Buscar "zi":   z → i → não é fim de palavra → prefixo existe
Buscar "abc":  nenhum filho 'a' na raiz → não encontrado

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

pub const TrieNode = struct {
    filhos: [26]?*TrieNode,
    fim_palavra: bool,
    contagem: u32, // quantas palavras passam por este nó

    pub fn init() TrieNode {
        return .{
            .filhos = [_]?*TrieNode{null} ** 26,
            .fim_palavra = false,
            .contagem = 0,
        };
    }
};

pub const Trie = struct {
    raiz: *TrieNode,
    allocator: Allocator,
    total_palavras: u32,

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

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

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

    /// Insere uma palavra na Trie.
    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) {
                const novo = try self.allocator.create(TrieNode);
                novo.* = TrieNode.init();
                atual.filhos[idx] = novo;
            }
            atual = atual.filhos[idx].?;
            atual.contagem += 1;
        }
        if (!atual.fim_palavra) {
            atual.fim_palavra = true;
            self.total_palavras += 1;
        }
    }

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

    /// Verifica se alguma palavra começa com o prefixo dado.
    pub fn temPrefixo(self: *const Trie, prefixo: []const u8) bool {
        return self.buscarNo(prefixo) != null;
    }

    /// Conta quantas palavras começam com o prefixo dado.
    pub fn contarPrefixo(self: *const Trie, prefixo: []const u8) u32 {
        const no = self.buscarNo(prefixo);
        return if (no) |n| n.contagem else 0;
    }

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

    /// Lista todas as palavras com o prefixo dado (autocompletar).
    pub fn autocompletar(
        self: *const Trie,
        allocator: Allocator,
        prefixo: []const u8,
        max_resultados: usize,
    ) ![][]u8 {
        var resultados = std.ArrayList([]u8).init(allocator);

        const no = self.buscarNo(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_resultados);

        return try resultados.toOwnedSlice();
    }

    fn coletarPalavras(
        no: *TrieNode,
        buffer: *std.ArrayList(u8),
        resultados: *std.ArrayList([]u8),
        allocator: Allocator,
        max: usize,
    ) !void {
        if (resultados.items.len >= max) return;

        if (no.fim_palavra) {
            const copia = try allocator.dupe(u8, buffer.items);
            try resultados.append(copia);
        }

        for (0..26) |i| {
            if (no.filhos[i]) |filho| {
                try buffer.append(@intCast('a' + i));
                try coletarPalavras(filho, buffer, resultados, allocator, max);
                _ = buffer.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();

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

    try stdout.print("Total de palavras: {d}\n\n", .{trie.total_palavras});

    // Busca
    try stdout.print("Buscar \"zig\": {}\n", .{trie.buscar("zig")});
    try stdout.print("Buscar \"zi\": {}\n", .{trie.buscar("zi")});
    try stdout.print("Prefixo \"zi\": {}\n", .{trie.temPrefixo("zi")});
    try stdout.print("Palavras com prefixo \"zi\": {d}\n\n", .{trie.contarPrefixo("zi")});

    // Autocompletar
    const sugestoes = try trie.autocompletar(allocator, "zi", 10);
    defer {
        for (sugestoes) |s| allocator.free(s);
        allocator.free(sugestoes);
    }

    try stdout.print("Autocompletar \"zi\":\n", .{});
    for (sugestoes) |s| {
        try stdout.print("  {s}\n", .{s});
    }
}