Estruturas de Dados 06 - Segment Tree++

Aula escrita por Leonardo Paes e Luca Dantas

Problemas envolvendo Segment Trees são muito recorrentes na OBI, principalmente na modalidade programação nível 2. Por exemplo, nas últimas 3 edições da olimpíada, OBI 2020, 2019 e 2018, houve problemas solucionáveis com Segment Trees.

Apesar de existirem diversos problemas cuja única dificuldade é saber implementar uma Segment Tree, problemas mais avançados requerem soluções que exijem mais esforço e criatividade. Nesses problemas, será necessário construir uma Segment Tree com um merge (ou combinação) específico, ao invés das usuais Segs de soma e máximo, por exemplo.

Mas o que é esse merge? É a função que recebe ambos os filhos de um nó da Seg como parâmetro e retorna como resultado os valores que o nó pai terá. Por exemplo: se precisarmos fazer uma seg de soma, a função seria semelhante a esta:

Problema de exemplo

Tendo em vista que problemas mais difíceis exigem merges mais complexos, vamos detalhar um passo-a-passo para resolvê-los por meio da explicação da solução do seguinte problema:

Problema: Dado um vetor $v$ de $n$ inteiros $v_i$, temos $q$ consultas a responder. Elas podem ser de dois tipos:

1. Mudar o valor do $i$-ésimo elemento de $v$ para $x$. Em outras palavras, $v_i$ recebe $x$.
2. Dado um intervalo $[l, r]$, encontrar o valor do subarray $[i, j]$ de soma máxima tal que $l \leq i \leq j \leq r$.

Restrições: $1 \leq n, q \leq 10^5, -10^9 \leq v_i \leq 10^9$

Link de possível judge para submissão do problema de exemplo (note que é necessário se registrar no EDU do CodeForces para acessar o link).

Um exemplo da consulta de tipo 2:

• Seja $v =$[$1, -5, 2, 5, -2$] e a consulta do tipo $2$, com intervalo [$1, 5]$. O valor que deverá ser imprimido é $7$, soma do intervalo contíguo $[3, 4]$, contido dentro do intervalo $[1, 5]$.

Para resolvermos problemas que envolvem merges complexos, é interessante o uso de uma Struct ou Class, para facilitar a escrita do código. Geralmente, criamos um Struct $Node$, que terá todas as variáveis necessárias para mantermos os valores desejados em cada nó da Seg. Por padrão, utilizaremos a variável $ans$, abreviação da palavra inglesa answer, que significa "resposta" em português. Se preferir, pode chamá-la de $res$. Note que estamos comentando os nomes que geralmente usamos para as variáveis.

Solução do problema

Segue abaixo um passo-a-passo de como resolver um problema de Segment Tree com merge complexo:

• Passo 1: Devemos estabelecer o valor de uma folha da Seg. Geralmente este passo é bem direto.
• Passo 2: Em seguida, temos que considerar a seguinte situação: Somente com o que está armazenado no filho da esquerda e da direita de qualquer nó que não é folha, deve ser possível conseguir a resposta para o próprio nó também. Ou seja, para deduzir o que armazenar nos nós, devemos pensar no que é necessário ter em ambos os filhos de cada nó para chegar na sua própria resposta.
• Passo 3: Após saber quais variáveis são necessários e suficientes, devemos entender como combiná-las corretamente.

Vamos aplicar o passo-a-passo no problema de exemplo:

Passo 1: Para o nó folha $(i, i), ans = max(v_i, 0)$. Note que o intervalo vazio sempre é uma possibilidade, por isso o valor $0$.

Passo 2: Com a variável $ans$ do filho da esquerda, sabemos a soma máxima dentre os subintervalos contíguos que começam e terminam no intervalo representado pelo nó do filho da esquerda. O mesmo vale para o filho da direita. Portanto, só nos resta calcular a soma máxima dentre os subintervalos contíguos que começam no filho da esquerda e terminam no filho da direita. Como esse subintervalo de soma máxima se parecerá? Ele certamente será formado por um sufixo do filho da esquerda concatenado a um prefixo do filho da direita. Não somente isso, será a concatenação do sufixo de soma máxima do filho da esquerda com o prefixo de soma máxima do filho da direita. Note que, como precisaremos dessas duas variáveis, precisamos mantê-las durante todos os nós. Portanto, precisaremos de uma variável que guarda a soma do nó também! Logo, conseguimos determinar as variáveis suficientes para manter em cada nó da Seg!

Passo 3: Como pode-se concluir a partir do passo 2, a resposta para um nó que não é folha é:

$max(filho_{esquerda}.ans, filho_{direita}.ans, filho_{esquerda}.suff + filho_{direita}.pref)$,

onde $.ans$, $.suff$ e $.pref$ guardam a resposta de um nó, o valor do sufixo e do prefixo de soma máxima, respectivamente. Além disso, o prefixo de soma máxima do nó atual é:

$max(filho_{esquerda}.pref, filho_{esquerda}.soma + filho_{direita}.pref$,

onde $.soma$ representa a soma total dos valores no intervalo de um nó. De modo análogo, o sufixo de soma máxima é

$max(filho_{direita}.suff, filho_{direita}.soma + filho_{esquerda}.suff$.

Segue abaixo um código com essa Seg para melhor entedimento (note que no código a função merge é chamada de join):

OBI 2020, Fase 3, P2: Candidatas

Agora vamos analisar um exemplo que caiu na OBI em 2020 e que utiliza essa técnica, o problema Candidatas.

Resumidamente, o problema nos dá um vetor e realizará dois tipos de operações:

1. Mudar o valor de uma posição do vetor;
2. Saber o número de subintervalos desse vetor com máximo divisor comum maior do que 1 dentro de um intervalo $[l, r]$ dado.

Como já estamos trabalhando técnicas de árvore em segmentos, pode-se inferir que esse problema têm uma solução envolvendo isso. Porém, podemos perceber que essa é uma boa técnica a ser utilizada neste problema pelo tipo de operações que nos são pedidas: mudar o valor em um posição e perguntar uma informação sobre um intervalo.

A base da solução para esse problema é bem parecida com a ideia geral apresentada na aula: queremos dividir as respostas para cada segmento pequeno e conseguir ter informação suficiente dentro deles para unir dois segmentos e calcular eficientemente assim a resposta de um segmento maior, assim como fazemos em problemas de árvore de segmentos normalmente. Portanto, vamos focar nossos esforços em responder o problema para um intervalo e em descobrir como calcular a resposta de um intervalo maior baseado na união de dois intervalos.

O caso base neste problema é um intervalo formado por um único elemento, e nesse caso a única resposta é formada pelo subintervalo formado somente por ele.

Para todos os conjunto maiores, as respostas têm 3 possibilidades:

1. A subsequência válida está completamente contida no filho da esquerda;
2. A subsequência válida está completamente contida no filho da direita;
3. A subsequência válida inicia no filho da esquerda e termina no filho da direita.

Para os dois primeiros casos basta adicionarmos na resposta do nó atual a soma das respostas dos dois filhos. Para o terceiro caso, já chegamos em um problema um pouco mais complicado.

Primeira ideia

Uma observação que podemos ter é que, nesta situação, a resposta vai ser formada por um prefixo dos elementos do filho da direita e um sufixo dos elementos do filho da esquerda. Uma ideia então seria salvarmos todos os sufixos e prefixos de cada segmento e, ao combiná-los, compararmos todos os prefixos da direita com os sufixos da esquerda um a um em $O(tamanho^2)$ e contarmos em quantos deles o MDC é maior que 1.

Vale lembrar que a operação de MDC é comutativa e associativa, então se quisermos calcular o MDC de um segmento e tivermos o MDC de dois subintervalos que juntos cobrem o segmento inteiro, o MDC do segmento inteiro vai ser simplesmente o MDC dos valores desses dois subsegmentos.

Otimizando o merge

Para otimizar a solução acima, vamos precisar fazer algumas observações sobre as propriedades do MDC.
Perceba que se analisarmos os prefixos (também vale para os sufixos pois eles são simétricos), a estrutura que encontraremos será uma sequência de valores decrescentes por uma propriedade dos MDCs. Podemos observá-la vendo que se tivermos um conjunto $A$ de elementos tal que $mdc(A) = x$ e adicionarmos um elemento $v$ ao conjunto, teremos que $mdc(A) = mdc(x, v) \leq x$. Por exemplo, cofira a imagem abaixo:

 

Outra observação que podemos fazer é que toda vez que o MDC é alterado, a decomposição em primos desse MDC "perde" um elemento. No caso acima, por exemplo, tínhamos: $2^2 \cdot 3 \cdot 5 \rightarrow 2^2 \cdot 3 \rightarrow 2 \cdot 3 \rightarrow 3 \rightarrow 3$. Isso acontece pois o MDC é essencialmente o "and" dos fatores primos de todos os elementos do conjunto, ou seja, para a potência de primo $p^k$ estar presente na fatoração do MDC do conjunto, todos os elementos do conjunto têm que ter essa potência de primo em sua fatoração.

Portanto, o número de MDC's diferentes para o prefixo é limitado por $log(max_v) = 30$, já que o maior valor de um elemento desse conjunto só pode ter no máximo 30 divisores primos e a cada vez que o valor do MDC muda um desses primos é perdido.

Juntando essas observações, podemos encontrar uma forma muito mais eficiente de salvar os valores de prefixos e sufixos. Já que somente existem $O(log(max))$ valores diferentes para o MDC, basta salvarmos cada intervalo de prefixos com o mesmo valor de MDC, salvando qual é esse valor e qual o tamanho desse conjunto. Isso pode ser armazenado facilmente em um vector<pair<int,int>>, indicando qual o valor e quantos itens existem com esse MDC.

Por exemplo, para a ilustração acima teríamos: $[(60, 1), (12, 1), (6, 1), (3, 2)]$.

Como calcular este vetor dos MDC's do prefixo e sufixo?

Para calcular o vetor de prefixo (o sufixo é simétrico, então a explicação será a mesma) do nó atual, iremos juntar duas partes. A primeira é simplesmente pegar o prefixo do filho da esquerda, pois essa parte se manterá inalterada. Para considerarmos os prefixos que terminam no filho da direita, vamos simplesmente adicionar os mesmos prefixos já salvos no filho da direita tirando o MDC do valor já presente nele com o MDC do filho da esquerda inteiro.

Como um exemplo, se o vetor de prefixo do filho da direita for $[(20, 1), (5, 2)]$ e o MDC do intervalo da esquerda inteiro era $10$, os valores que iremos adicionar serão $[(10, 1), (5, 2)]$. Porém, tome cuidado, pois se somente adicionarmos assim, poderíamos deixar segmentos com o mesmo valor separado. Imagine que no caso mostrado agora o MDC do nó da esquerda é $5$. Nesse caso, se só seguirmos cegamente o algoritmo de adicionar os valores tirando o MDC com o nó da esquerda completo, adicionaríamos $[(5, 1), (5, 2)]$, o que quebra a propriedade de que mantemos os conjuntos de mesmo valor em um mesmo elemento do vetor, fazendo com que o vetor inteiro possa ter mais que $O(log(max))$ elementos. Para remediar esse problema, basta toda vez que formos adicionar um elemento, checarmos se o novo MDC que estamos adicionando é igual ao último elemento do vetor atualmente; caso seja, somente somamos os tamanhos do intervalo ao invés de adicionar outro elemento ao vetor.

Um outro detalhe que utilizamos para tornar a implementação mais simples é salvar apenas os prefixos cujo MDC é maior do que 1.

Segue uma implementação dessa ideia para melhor compreensão:

Como podemos calcular a resposta com esses valores de prefixo e sufixo?

A primeira solução (e menos eficiente) é simplesmente repetir o que fizemos na primeira ideia e compararmos os pares um a um em $O(tamanho^2)$. Nesse caso, como o tamanho é $O(log(max))$, o merge possuirá complexidade $O(log^2)$ e a solução completa $O(n log^3)$, o que não é suficiente para a solução final.

A observação final é que se temos um prefixo $p_i$ e um sufixo do outro nó $s_j$ e $MDC(p_i, s_j) > 1$, o $MDC(p_i, s_k) > 1$ para qualquer $k \leq j$. De maneira semelhante, se exise um $j$ tal que $MDC(p_i, s_j) = 1$, $MDC(p_i, s_k) = 1$ para qualquer $k \geq j$. Note que a análise é a mesma se trocarmos a posição de prefixo e sufixo aqui pois eles são simétricos.

Em outras palavras, existe um ponto $r$ em que todos os intervalos iniciados em $p_i$ e que terminam até $r$ são maiores que $1$, e portanto são válidos para a resposta. E todos os intervalos que terminam depois de $r$ não são válidos para a resposta. Logo, podemos fazer uma busca mais eficiente para achar quantos valores iniciando em $p_i$ são válidos.

Essa busca mais eficiente pode ser uma busca binária, levando a complexidade à $O(n log^2(n) \cdot log(log(n)))$, o que é uma melhoria. Porém, a implementação se torna complicada considerando como fizemos o cálculo dos MDC's de prefixos e sufixos acima. Além disso, esta solução ainda pode não ser rápida o suficiente. Portanto, implementaremos a técnica de Two Pointers.

Para perceber que Two Pointers funciona nesse caso, basta perceber que os pontos ótimos dos prefixos $p_i$ nunca aumentarão se colocarmos mais elementos ao conjunto. Considerando que o ponto ótimo do prefixo $p_i$ é $o_i$, temos que para todo $i < j, o_i \geq o_j$. Ou seja, os pontos ótimos são decrescentes, portanto aplicar Two Pointers é válido.

Segue uma implementação desse cálculo da resposta para mais detalhes (os detalhes de como implementar estão comentados no código):

Segue o código agora da solução completa: