Segment Tree em Zig — Implementação Completa

Segment Tree em Zig — Implementação Completa

A Segment Tree (árvore de segmentos) permite consultas e atualizações em intervalos de um array em O(log n). Suporta operações como soma, mínimo, máximo e GCD sobre faixas.

Conceito

Array: [1, 3, 5, 7, 9, 11]

Segment Tree (soma):
                  [0,5]=36
                /         \
           [0,2]=9       [3,5]=27
           /    \        /     \
       [0,1]=4 [2,2]=5 [3,4]=16 [5,5]=11
       /   \           /    \
   [0,0]=1 [1,1]=3 [3,3]=7 [4,4]=9

Consulta soma(1, 4):
  = [1,1] + [2,2] + [3,4]
  = 3 + 5 + 16 = 24

Atualizar arr[2] = 10:
  Atualiza [2,2]=10, [0,2]=14, [0,5]=41

Implementação em Zig

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

/// Segment Tree genérica para consultas em intervalos.
pub fn SegmentTree(comptime T: type) type {
    return struct {
        const Self = @This();

        arvore: []T,
        n: usize,
        neutro: T,
        combinar: *const fn (T, T) T,
        allocator: Allocator,

        /// Inicializa a segment tree com um array de dados.
        pub fn init(
            allocator: Allocator,
            dados: []const T,
            neutro: T,
            combinar: *const fn (T, T) T,
        ) !Self {
            const n = dados.len;
            const arvore = try allocator.alloc(T, 4 * n);
            @memset(arvore, neutro);

            var st = Self{
                .arvore = arvore,
                .n = n,
                .neutro = neutro,
                .combinar = combinar,
                .allocator = allocator,
            };

            st.construir(dados, 1, 0, n - 1);
            return st;
        }

        pub fn deinit(self: *Self) void {
            self.allocator.free(self.arvore);
        }

        fn construir(self: *Self, dados: []const T, pos: usize, ini: usize, fim: usize) void {
            if (ini == fim) {
                self.arvore[pos] = dados[ini];
                return;
            }
            const meio = ini + (fim - ini) / 2;
            self.construir(dados, 2 * pos, ini, meio);
            self.construir(dados, 2 * pos + 1, meio + 1, fim);
            self.arvore[pos] = self.combinar(self.arvore[2 * pos], self.arvore[2 * pos + 1]);
        }

        /// Consulta o resultado da operação no intervalo [l, r] — O(log n).
        pub fn consultar(self: *const Self, l: usize, r: usize) T {
            return self.consultarHelper(1, 0, self.n - 1, l, r);
        }

        fn consultarHelper(self: *const Self, pos: usize, ini: usize, fim: usize, l: usize, r: usize) T {
            if (r < ini or fim < l) return self.neutro;
            if (l <= ini and fim <= r) return self.arvore[pos];
            const meio = ini + (fim - ini) / 2;
            return self.combinar(
                self.consultarHelper(2 * pos, ini, meio, l, r),
                self.consultarHelper(2 * pos + 1, meio + 1, fim, l, r),
            );
        }

        /// Atualiza o valor na posição idx — O(log n).
        pub fn atualizar(self: *Self, idx: usize, valor: T) void {
            self.atualizarHelper(1, 0, self.n - 1, idx, valor);
        }

        fn atualizarHelper(self: *Self, pos: usize, ini: usize, fim: usize, idx: usize, valor: T) void {
            if (ini == fim) {
                self.arvore[pos] = valor;
                return;
            }
            const meio = ini + (fim - ini) / 2;
            if (idx <= meio) {
                self.atualizarHelper(2 * pos, ini, meio, idx, valor);
            } else {
                self.atualizarHelper(2 * pos + 1, meio + 1, fim, idx, valor);
            }
            self.arvore[pos] = self.combinar(self.arvore[2 * pos], self.arvore[2 * pos + 1]);
        }
    };
}

fn somaI64(a: i64, b: i64) i64 {
    return a + b;
}

fn minimoI64(a: i64, b: i64) i64 {
    return @min(a, b);
}

fn maximoI64(a: i64, b: i64) i64 {
    return @max(a, b);
}

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

    const dados = [_]i64{ 1, 3, 5, 7, 9, 11 };

    // Segment Tree de soma
    var st_soma = try SegmentTree(i64).init(allocator, &dados, 0, &somaI64);
    defer st_soma.deinit();

    try stdout.print("=== Segment Tree (Soma) ===\n", .{});
    try stdout.print("Soma [0,5] = {d}\n", .{st_soma.consultar(0, 5)});
    try stdout.print("Soma [1,4] = {d}\n", .{st_soma.consultar(1, 4)});
    try stdout.print("Soma [2,3] = {d}\n", .{st_soma.consultar(2, 3)});

    st_soma.atualizar(2, 10);
    try stdout.print("Após arr[2]=10: Soma [0,5] = {d}\n", .{st_soma.consultar(0, 5)});

    // Segment Tree de mínimo
    var st_min = try SegmentTree(i64).init(allocator, &dados, std.math.maxInt(i64), &minimoI64);
    defer st_min.deinit();

    try stdout.print("\n=== Segment Tree (Mínimo) ===\n", .{});
    try stdout.print("Min [0,5] = {d}\n", .{st_min.consultar(0, 5)});
    try stdout.print("Min [2,4] = {d}\n", .{st_min.consultar(2, 4)});
}

Análise de Complexidade

OperaçãoTempoEspaço
ConstruirO(n)O(4n)
ConsultarO(log n)O(log n)
AtualizarO(log n)O(log n)

Recursos Relacionados

Continue aprendendo Zig

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