Segment Tree em Rust

Implemente Segment Tree em Rust para consultas de intervalo e atualizações pontuais. Lazy propagation, diagramas visuais e Big-O.

A Segment Tree (Árvore de Segmentos) é uma estrutura de dados que permite realizar consultas de intervalo (range queries) e atualizações de forma eficiente. Dada uma sequência de n elementos, a Segment Tree responde a perguntas como “qual a soma dos elementos de i a j?” ou “qual o mínimo no intervalo [i, j]?” em O(log n), mesmo quando os elementos são modificados dinamicamente.

Essa estrutura é amplamente utilizada em programação competitiva e em aplicações que exigem consultas rápidas sobre intervalos de dados que mudam frequentemente, como sistemas de monitoramento, bancos de dados analíticos e engines de jogos (detecção de colisão em intervalos).

Como Funciona

A Segment Tree é uma árvore binária completa onde cada nó representa um intervalo do array original. A raiz representa o array inteiro, e cada folha representa um único elemento. Nós internos armazenam o resultado agregado (soma, mínimo, máximo, etc.) de seus filhos.

Array original: [2, 1, 5, 3, 4, 7]
Operação: soma de intervalo

Árvore de Segmentos (cada nó mostra [intervalo]: valor):

                    [0..5]: 22
                   /        \
            [0..2]: 8        [3..5]: 14
           /      \          /       \
      [0..1]: 3  [2]: 5  [3..4]: 7  [5]: 7
      /    \              /    \
  [0]: 2  [1]: 1     [3]: 3  [4]: 4

Consulta: soma(1, 4) = ?
  Decompõe em: [1] + [2] + [3..4]
             =  1  +  5  +   7   = 13

Passo a passo:
  [0..5] → precisa descer (intervalo parcial)
    [0..2] → parcial
      [0..1] → parcial
        [0] → fora do intervalo, ignora
        [1] → dentro! retorna 1
      [2] → dentro! retorna 5
    [3..5] → parcial
      [3..4] → totalmente dentro! retorna 7
      [5] → fora do intervalo, ignora
  Resultado: 1 + 5 + 7 = 13

Atualização de um elemento:

Atualizar: array[2] = 10 (era 5, diferença = +5)

Caminho de atualização (baixo para cima):
  [2]: 5 → 10
  [0..2]: 8 → 13
  [0..5]: 22 → 27

Árvore após atualização:

                    [0..5]: 27
                   /        \
           [0..2]: 13        [3..5]: 14
           /      \          /       \
      [0..1]: 3  [2]: 10 [3..4]: 7  [5]: 7

A árvore é armazenada em um array linear, onde para o nó no índice i:

  • Filho esquerdo: 2*i + 1
  • Filho direito: 2*i + 2
  • Pai: (i - 1) / 2
Representação em array (indexação 0-based):

Índice:   0     1     2     3     4     5     6     7     8     9    10    11
Valor:  [22]  [ 8]  [14]  [ 3]  [ 5]  [ 7]  [ 7]  [ 2]  [ 1]  [ 3]  [ 4]  [--]
         raiz  esq   dir  ...folhas do array original nos índices 7..12

Complexidade

OperaçãoTempoEspaço
ConstruçãoO(n)O(4n) = O(n)
Consulta de intervaloO(log n)O(log n) (pilha)
Atualização pontualO(log n)O(log n) (pilha)
Atualização de intervalo (lazy)O(log n)O(log n) (pilha)
  • Construção O(n): cada elemento é processado uma vez, e os nós internos são computados bottom-up.
  • Consulta e atualização O(log n): no máximo O(log n) nós são visitados em cada nível da árvore.
  • Espaço O(4n): no pior caso, a árvore precisa de até 4n posições no array (para n que não é potência de 2).

Implementação em Rust

/// Segment Tree para consultas de soma em intervalos com atualização pontual.
struct SegmentTree {
    n: usize,
    arvore: Vec<i64>,
}

impl SegmentTree {
    /// Constrói a Segment Tree a partir de um slice de dados.
    fn construir(dados: &[i64]) -> Self {
        let n = dados.len();
        let mut arvore = vec![0i64; 4 * n];
        if n > 0 {
            Self::build(&mut arvore, dados, 0, 0, n - 1);
        }
        SegmentTree { n, arvore }
    }

    fn build(arvore: &mut [i64], dados: &[i64], no: usize, inicio: usize, fim: usize) {
        if inicio == fim {
            arvore[no] = dados[inicio];
            return;
        }
        let meio = (inicio + fim) / 2;
        Self::build(arvore, dados, 2 * no + 1, inicio, meio);
        Self::build(arvore, dados, 2 * no + 2, meio + 1, fim);
        arvore[no] = arvore[2 * no + 1] + arvore[2 * no + 2];
    }

    /// Consulta a soma no intervalo [esq, dir] (inclusive).
    fn consultar(&self, esq: usize, dir: usize) -> i64 {
        self.query(&self.arvore, 0, 0, self.n - 1, esq, dir)
    }

    fn query(&self, arvore: &[i64], no: usize, inicio: usize, fim: usize,
             esq: usize, dir: usize) -> i64 {
        // Intervalo totalmente fora
        if dir < inicio || fim < esq {
            return 0;
        }
        // Intervalo totalmente dentro
        if esq <= inicio && fim <= dir {
            return arvore[no];
        }
        // Intervalo parcial — desce para os filhos
        let meio = (inicio + fim) / 2;
        let soma_esq = self.query(arvore, 2 * no + 1, inicio, meio, esq, dir);
        let soma_dir = self.query(arvore, 2 * no + 2, meio + 1, fim, esq, dir);
        soma_esq + soma_dir
    }

    /// Atualiza o valor na posição idx para novo_valor.
    fn atualizar(&mut self, idx: usize, novo_valor: i64) {
        Self::update(&mut self.arvore, 0, 0, self.n - 1, idx, novo_valor);
    }

    fn update(arvore: &mut [i64], no: usize, inicio: usize, fim: usize,
              idx: usize, valor: i64) {
        if inicio == fim {
            arvore[no] = valor;
            return;
        }
        let meio = (inicio + fim) / 2;
        if idx <= meio {
            Self::update(arvore, 2 * no + 1, inicio, meio, idx, valor);
        } else {
            Self::update(arvore, 2 * no + 2, meio + 1, fim, idx, valor);
        }
        arvore[no] = arvore[2 * no + 1] + arvore[2 * no + 2];
    }
}

fn main() {
    let dados = vec![2, 1, 5, 3, 4, 7];
    let mut st = SegmentTree::construir(&dados);

    println!("Array: {:?}", dados);
    println!("Soma [0, 5]: {}", st.consultar(0, 5)); // 22
    println!("Soma [1, 4]: {}", st.consultar(1, 4)); // 13
    println!("Soma [2, 2]: {}", st.consultar(2, 2)); // 5

    // Atualizar posição 2: 5 → 10
    st.atualizar(2, 10);
    println!("\nApós atualizar pos 2 para 10:");
    println!("Soma [0, 5]: {}", st.consultar(0, 5)); // 27
    println!("Soma [1, 4]: {}", st.consultar(1, 4)); // 18

    // Atualizar posição 0: 2 → 8
    st.atualizar(0, 8);
    println!("\nApós atualizar pos 0 para 8:");
    println!("Soma [0, 5]: {}", st.consultar(0, 5)); // 33
    println!("Soma [0, 0]: {}", st.consultar(0, 0)); // 8
}

Exemplo de Execução

Array: [2, 1, 5, 3, 4, 7]
Soma [0, 5]: 22
Soma [1, 4]: 13
Soma [2, 2]: 5

Após atualizar pos 2 para 10:
Soma [0, 5]: 27
Soma [1, 4]: 18

Após atualizar pos 0 para 8:
Soma [0, 5]: 33
Soma [0, 0]: 8

Trace da consulta soma(1, 4) no array [2, 1, 5, 3, 4, 7]:

query(no=0, [0..5], consulta [1..4])
  → parcial, desce
  ├─ query(no=1, [0..2], consulta [1..4])
  │    → parcial, desce
  │    ├─ query(no=3, [0..1], consulta [1..4])
  │    │    → parcial, desce
  │    │    ├─ query(no=7, [0..0], consulta [1..4])
  │    │    │    → fora! retorna 0
  │    │    └─ query(no=8, [1..1], consulta [1..4])
  │    │         → dentro! retorna 1
  │    │    retorna 0 + 1 = 1
  │    └─ query(no=4, [2..2], consulta [1..4])
  │         → dentro! retorna 5
  │    retorna 1 + 5 = 6
  └─ query(no=2, [3..5], consulta [1..4])
       → parcial, desce
       ├─ query(no=5, [3..4], consulta [1..4])
       │    → totalmente dentro! retorna 7
       └─ query(no=6, [5..5], consulta [1..4])
            → fora! retorna 0
       retorna 7 + 0 = 7
  retorna 6 + 7 = 13

Variações e Otimizações

1. Segment Tree para Mínimo de Intervalo (RMQ)

/// Segment Tree para Range Minimum Query.
struct SegTreeMin {
    n: usize,
    arvore: Vec<i64>,
}

impl SegTreeMin {
    fn construir(dados: &[i64]) -> Self {
        let n = dados.len();
        let mut arvore = vec![i64::MAX; 4 * n];
        if n > 0 {
            Self::build(&mut arvore, dados, 0, 0, n - 1);
        }
        SegTreeMin { n, arvore }
    }

    fn build(arvore: &mut [i64], dados: &[i64], no: usize, ini: usize, fim: usize) {
        if ini == fim {
            arvore[no] = dados[ini];
            return;
        }
        let meio = (ini + fim) / 2;
        Self::build(arvore, dados, 2 * no + 1, ini, meio);
        Self::build(arvore, dados, 2 * no + 2, meio + 1, fim);
        arvore[no] = arvore[2 * no + 1].min(arvore[2 * no + 2]);
    }

    fn consultar_min(&self, esq: usize, dir: usize) -> i64 {
        self.query(0, 0, self.n - 1, esq, dir)
    }

    fn query(&self, no: usize, ini: usize, fim: usize, esq: usize, dir: usize) -> i64 {
        if dir < ini || fim < esq {
            return i64::MAX;
        }
        if esq <= ini && fim <= dir {
            return self.arvore[no];
        }
        let meio = (ini + fim) / 2;
        let min_esq = self.query(2 * no + 1, ini, meio, esq, dir);
        let min_dir = self.query(2 * no + 2, meio + 1, fim, esq, dir);
        min_esq.min(min_dir)
    }

    fn atualizar(&mut self, idx: usize, valor: i64) {
        Self::update(&mut self.arvore, 0, 0, self.n - 1, idx, valor);
    }

    fn update(arvore: &mut [i64], no: usize, ini: usize, fim: usize, idx: usize, val: i64) {
        if ini == fim {
            arvore[no] = val;
            return;
        }
        let meio = (ini + fim) / 2;
        if idx <= meio {
            Self::update(arvore, 2 * no + 1, ini, meio, idx, val);
        } else {
            Self::update(arvore, 2 * no + 2, meio + 1, fim, idx, val);
        }
        arvore[no] = arvore[2 * no + 1].min(arvore[2 * no + 2]);
    }
}

fn main() {
    let dados = vec![5, 2, 8, 1, 4, 7, 3, 6];
    let mut st = SegTreeMin::construir(&dados);

    println!("Array: {:?}", dados);
    println!("Min [0, 7]: {}", st.consultar_min(0, 7)); // 1
    println!("Min [0, 2]: {}", st.consultar_min(0, 2)); // 2
    println!("Min [4, 7]: {}", st.consultar_min(4, 7)); // 3

    st.atualizar(3, 9); // muda 1 para 9
    println!("\nApós mudar pos 3 para 9:");
    println!("Min [0, 7]: {}", st.consultar_min(0, 7)); // 2
}

2. Lazy Propagation para Atualizações de Intervalo

Quando precisamos atualizar todos os elementos em um intervalo (ex: somar +5 a todos no intervalo [2, 6]), Lazy Propagation evita atualizar cada elemento individualmente.

/// Segment Tree com Lazy Propagation para atualização de intervalo.
struct SegTreeLazy {
    n: usize,
    arvore: Vec<i64>,
    lazy: Vec<i64>,  // atualizações pendentes
}

impl SegTreeLazy {
    fn construir(dados: &[i64]) -> Self {
        let n = dados.len();
        let mut arvore = vec![0i64; 4 * n];
        let lazy = vec![0i64; 4 * n];
        if n > 0 {
            Self::build(&mut arvore, dados, 0, 0, n - 1);
        }
        SegTreeLazy { n, arvore, lazy }
    }

    fn build(arvore: &mut [i64], dados: &[i64], no: usize, ini: usize, fim: usize) {
        if ini == fim {
            arvore[no] = dados[ini];
            return;
        }
        let meio = (ini + fim) / 2;
        Self::build(arvore, dados, 2 * no + 1, ini, meio);
        Self::build(arvore, dados, 2 * no + 2, meio + 1, fim);
        arvore[no] = arvore[2 * no + 1] + arvore[2 * no + 2];
    }

    /// Propaga atualizações pendentes para os filhos.
    fn propagar(&mut self, no: usize, ini: usize, fim: usize) {
        if self.lazy[no] != 0 {
            let meio = (ini + fim) / 2;
            // Aplica a atualização pendente ao nó atual
            self.arvore[no] += self.lazy[no] * (fim - ini + 1) as i64;

            // Se não é folha, propaga para os filhos
            if ini != fim {
                self.lazy[2 * no + 1] += self.lazy[no];
                self.lazy[2 * no + 2] += self.lazy[no];
            }
            self.lazy[no] = 0;
        }
    }

    /// Soma `valor` a todos os elementos no intervalo [esq, dir].
    fn atualizar_intervalo(&mut self, esq: usize, dir: usize, valor: i64) {
        self.update_range(0, 0, self.n - 1, esq, dir, valor);
    }

    fn update_range(&mut self, no: usize, ini: usize, fim: usize,
                    esq: usize, dir: usize, valor: i64) {
        self.propagar(no, ini, fim);

        if dir < ini || fim < esq {
            return;
        }
        if esq <= ini && fim <= dir {
            self.lazy[no] += valor;
            self.propagar(no, ini, fim);
            return;
        }
        let meio = (ini + fim) / 2;
        self.update_range(2 * no + 1, ini, meio, esq, dir, valor);
        self.update_range(2 * no + 2, meio + 1, fim, esq, dir, valor);
        self.arvore[no] = self.arvore[2 * no + 1] + self.arvore[2 * no + 2];
    }

    /// Consulta a soma no intervalo [esq, dir].
    fn consultar(&mut self, esq: usize, dir: usize) -> i64 {
        self.query_range(0, 0, self.n - 1, esq, dir)
    }

    fn query_range(&mut self, no: usize, ini: usize, fim: usize,
                   esq: usize, dir: usize) -> i64 {
        self.propagar(no, ini, fim);

        if dir < ini || fim < esq {
            return 0;
        }
        if esq <= ini && fim <= dir {
            return self.arvore[no];
        }
        let meio = (ini + fim) / 2;
        let s1 = self.query_range(2 * no + 1, ini, meio, esq, dir);
        let s2 = self.query_range(2 * no + 2, meio + 1, fim, esq, dir);
        s1 + s2
    }
}

fn main() {
    let dados = vec![1, 3, 5, 7, 9, 11];
    let mut st = SegTreeLazy::construir(&dados);

    println!("Array: {:?}", dados);
    println!("Soma [1, 3]: {}", st.consultar(1, 3)); // 15

    // Somar +10 a todos no intervalo [1, 4]
    st.atualizar_intervalo(1, 4, 10);
    println!("\nApós somar +10 no intervalo [1,4]:");
    println!("Soma [1, 3]: {}", st.consultar(1, 3)); // 45 (13+15+17)
    println!("Soma [0, 5]: {}", st.consultar(0, 5)); // 76

    // Somar +5 a todos no intervalo [0, 2]
    st.atualizar_intervalo(0, 2, 5);
    println!("\nApós somar +5 no intervalo [0,2]:");
    println!("Soma [0, 5]: {}", st.consultar(0, 5)); // 91
}

Aplicações Práticas

  • Consultas de intervalo em bancos de dados: somar vendas em um período, encontrar a temperatura mínima em uma faixa de dias.
  • Compressão de coordenadas: em geometria computacional, mapear coordenadas para intervalos e consultar propriedades.
  • Problemas de programação competitiva: inúmeros problemas de competição envolvem range queries com atualizações dinâmicas.
  • Processamento de sinais: calcular médias móveis e estatísticas sobre janelas deslizantes.
  • Jogos: detectar colisão em intervalos, calcular dano total sobre faixas de personagens.
  • Finanças: consultar o preço mínimo/máximo de ações em intervalos de tempo.

Comparação com Alternativas

CaracterísticaSegment TreeFenwick Tree (BIT)Sparse TablePrefix Sum
Consulta de intervaloO(log n)O(log n)O(1)O(1)
Atualização pontualO(log n)O(log n)O(n)O(n)
Atualização de intervaloO(log n)*O(log n)*ImpossívelImpossível
EspaçoO(4n)O(n)O(n log n)O(n)
ImplementaçãoMédiaSimplesSimplesTrivial
FlexibilidadeAlta (min, max, soma, gcd…)Limitada (soma, xor)Somente consultaSomente soma

*Com lazy propagation. A Segment Tree e a mais flexível, suportando qualquer operação associativa. Use Fenwick Tree quando só precisa de soma/xor e quer implementação mais simples.

Exercícios Práticos

  1. Range GCD: Modifique a Segment Tree para calcular o MDC (Greatest Common Divisor) de um intervalo. Dica: substitua a soma por gcd(a, b) usando o algoritmo de Euclides.

  2. Contagem em intervalo: Construa uma Segment Tree que conta quantos elementos no intervalo [l, r] são maiores que um valor k. Dica: use um merge sort tree (cada nó armazena os elementos ordenados do intervalo).

  3. Inversões com Segment Tree: Dado um array, conte o número de inversões (pares i < j onde a[i] > a[j]) usando uma Segment Tree de frequência. Processe o array da direita para a esquerda.

  4. Segment Tree persistente: Implemente uma versão persistente da Segment Tree que preserva versões anteriores após atualizações. Cada atualização cria O(log n) novos nós, mantendo as versões antigas intactas.

Veja Também