Segment Tree

Range Sum Query usando uma Segment Tree

A ideia aqui é resolver o problema RSQ em que são necessárias atualizações nos elementos do array e também muitas consultas são feitas. Usando uma Segment Tree (ST), teremos as seguintes etapas:

• Construção da árvore a partir do array de números: O(n)
• Consultas (rsq): O(log(n))
• Propagar uma modificação do array na árvore: O(log(n))

Construção de uma ST

A maneira mais fácil de explicar uma ST é através de um exemplo. Suponha que os elementos do array são:

5 2 8 9 3 6

Nesse caso, temos n=6.

Iremos armazenar os elementos em um vector e ignoraremos a primeira posição, só para facilitar a construção e navegação da árvore. Assim, o primeiro elemento será o elemento de índice 1:

a = [-1, 5, 2, 8, 9, 3, 6]

A ideia da ST é contruir quebrar o array em vários intervalos. Então, com base nesses intervalos, construimos uma árvore onde cada nó guarda a soma dos elementos dos intervalos.

Vamos adotar uma estratégia de dividir e conquistar, sempre quebrando o intervalo atual ao meio. O primeiro intervalo é o array inteiro, representado pelos índices: [1, 7[ , onde 1 é o índice 1 do array (intervalo fechado, ou seja, o elemento do índice 1 está incluido no intervalo) e 7 é o último índice daquele intervalo, mas nesse caso o intervalo é aberto, ou seja, o elemento do índice 7 não está incluído, sendo portanto, o elemento de índice 6 o último incluído.

Esse nó vai armazenar a soma do intervalo [1,7[. Mas ainda não temos o valor dessa soma. Vamos dividir esse intervalo em 2:

• Esquerda: de 1 até a metade. 1 até 1+7/2 = [1, 1+ 3.5[ = [1 + 4.5[ = [1, 4[ (sempre arredondaremos para baixo)
• Direta: da metade até o final do intervalo. [4, 7[

Para cada nó filho, iremos repetir esse procedimento, sempre quebrando o intervalo em dois. Logo, para o nó [1,4[, quebraremos em dois intervalos: [1,2[ e [2,4[

Faremos isso até que o intervalo contenha apenas um elemento. Nesse exemplo, o intervalo [1, 2[ contém apenas um elemento, que é o próprio um (lembre que o intervalo é aberto em 2). Como cada nó guarda o valor da soma do intervalo, o valor da soma de um intervalo com um só elemento é o próprio elemento. Logo, o valor do nó [1,2[ será o valor do elemento da posição 1, nesse caso: 5. E a nossa árvore ficará assim:

Ainda não discutimos como iremos representar essa árvore. Você pode usar ponteiros se quiser, mas eu vou utilizar arrays. Sendo assim, o nó raiz ficará na posição 1, e os seus filhos na posição 2*1 e 2*1 +1. Assim, um nó x, terá filhos nas posições 2*x e 2*x +1.

Agora nossa árvore ficará assim:

Como se trata de uma RSQ, os nós guardam os valores das somas dos intervalos. Logo, como estamos dividindo o intervalo em 2, para obtermos o valor da some de um intervalo que foi dividido, basta somar os dois intervalos. Por exemplo, o nó [2,4[ guardará o valor do nó [2,3[ + o valor do nó [3,4[.

Sendo assim, agora podemos atualizar o valor dos nós intermediários:

Prontinho, nossa árvore está construída.

Agora falta discutir como iremos usá-la para fazer as consultas.

Consulta em uma ST

Suponha que a nossa consulta seja: rsq(1,7).

Ou seja, queremos a soma de todos os elementos das posições 1 a 6.

Essa seria o melhor caso de consulta. Pois o nó raiz da nossa árvore guarda exatamente essa informação. Portanto, basta retornar o valor do nó raiz.

Suponha que a consulta seja: rsq(1,3). Nesse caso, estamos interessados no valor da soma dos elementos { a[1]. a[2] }

O nó raiz, contém a soma de um intervalo maior do que o que procuramos, portanto, precisamos descer mais na árvore.

Vamos examinar o nó esquerdo primeiro: [1,4[ . O intervalo que esse nó representa é dos elementos { a[1], a[2], a[3] }, portanto, também é maior do que o que desejamos. Vamos descer mais um nível, também começando pelo nó esquerdo.

Chegamos no nó [1, 2[ . Esse nó guarda o valor da soma do conjunto: { a[1] } . Dessa vez, o intervalo que encontramos é menor que o intervalo procurado. E além disso, o intervalo encontrado está contido no intervalo procurado. Ou seja, esse valor encontrado precisará ser usado e somado com os demais valores encontrados. Assim, vamos utilizar esse valor. E vamos explorar o outro filho: [2,4[.

Esse outro filho guarda o valor do intervalo { a[2], a[3] }. Só para lembrar, o interavlo procurado é {a[1], a[2]} . Assim, o intervalo desse nó, contém valores que não nos interessam (a[3]), mas ao mesmo tempo contém valores que interessam (a[2]). Sendo assim, continuamos descendo na árvore.

Chegamos no nó [2, 3[ . Esse nó contém o valor do intervalo { a[2] }. Logo, esse valor nos interessa e vamos retornar o valor desse nó. Ao retornar, vamos analisar o outro filho:

O nó [3,4[ guarda o intervalo { a[3] } . Esse valor está fora do intervalo que queremos saber, portanto não nos interessa. Assim, voltamos a subir na árvore retornando os valores que nos interessam:

Assim, retornamos ao nó raiz. Ainda falta avaliar o filho direito. Esse filho direito guarda o intervalo { a[4], a[5], a[6] }. Ou seja, todos os valores do filho direito não nos interessam. O intervalo do filho direito está completamente fora do interavalo procurado. Sendo assim, podemos já descartar esse filho, não sendo necessário descer na árvore.

E assim, temos o valor do rsq(1,3) = 7.

Essa estrutura tem o seu pior caso para a consulta rsq da ordem de O (log(n)).

Agora veja o código abaixo com detalhes para ver como a construção e a consulta ocorrem na prática.

/*

Segment Tree

Rodrigo Paes

RSQ: Range SUM Queries: soma de intervalos com atualização dinâmica.

Caso os valores do array não precisem ser atualizados, utilizar

sparse table.

Essa árvore possui implementação baseada em: http://codeforces.com/blog/entry/15729

*/

#include <bits/stdc++.h>

using namespace std;

typedef vector<int> vi;

#define pb push_back

vi st; //segment tree

vi a; // list of elements

/*

Constrói a segment tree com base no vector a.

O(n)

*/

void build(int idx=1, int left=1, int right=a.size())

{   

    if (right - left < 2) // se o intervalo tiver 1 elemento, ou zero

    {

        st[idx] = a[left];

        return;

    }

    int mid = (left+right)  / 2;

    build(idx * 2, left, mid);

    build(idx * 2+1, mid, right);

    st[idx] = st[idx*2] + st[idx*2+1];

}

/*

Em uma segment tree, cada nó da árvore guarda o valor da soma de um determinado

intervalo.

Ou seja, cada posição idx da árvore guarda o valor da soma do intervalo [left, right[

(fechado em left, aberto em right).

x = posição inicial do intervalo a ser somado (intervalo fechado)

y = posição final do intervalo a ser somado (intervalo aberto)

idx = índice atual da segment tree

left = início do intervalo da soma do nó idx

right = fim do intervalo da soma do nó idx

O(log(n))

*/

int sum(int x, int y, int idx = 1, int left = 1, int right = a.size())

{   

    if ( x >= right || y <= left)

    {

        /* se o início do intervalo desejado (x) for maior ou igual ao right

        significa que o intervalo está à direita desse nó idx e, portanto,

        fora do intervalo desse nó.

        O mesmo raciocínio se aplica ao fim. Ou seja, se o fim do intervalo

        é menor que o início do intervalo do nó idx significa que o intervalo

        está à esquerda do nó.

        */

        return 0;

    }

    if ( x <= left  && y >= right)

    {

        /*

        Nesse caso temos:

            left      right

              |            |

              |                |

              x                y

        O intervalo do nó está contido no intervalo desejado, assim,

        não faz sentido dividir o intervalo em pedaços ainda menores,

        ou seja, não faz sentido continuar descendo na árvore ainda mais

        pois a informação mais agregada já está contida nesse nó

        */

        return st[idx];

    }

    /*

    Se chegamos aqui, significa que o intervalo do nó contém a soma de elementos

    que não são desejados pelo inervalo x,y. Ou seja, pode ser que

    [x,y[ seja menor que [left,right[, o que indica que [left,right[ tem um valor maior

    do que o necessário. Nesse caso, temos que descer na árvore até encontrar

    os intervalos que contenham exatamanete a informação que precisamos.

    */   

    int mid = (left+right) / 2;   

    return     sum (x,y, idx*2, left, mid) +

            sum(x, y, idx*2+1, mid, right);

}

/*

Caso uma posição do array tenha o seu valor modificado, ao invés de

modificar diretamente no array, deve-se chamar essa função modify

para alterar também a árvore.

Percebe que a alteração deve ser propagada por todos os nós da árvore

que possuem aquele valor somado.

O(log(n))

*/

void modify(int pos, int value, int idx = 1, int left=1, int right = a.size())

{   

    /* modifica apenas o valor relativo ao anterior.

    Por exemplo, se antes a[pos] fosse 10 e agora é 5, então temos que subtrair 5

    das somas.

    5 - 10 = -5

    */

    st[idx] += ( value - a[pos] );

    if (right - left < 2)

    {

        a[pos] = value;

        return;

    }

    int mid = (left+right) / 2;

    // Só precisa alterar um ramo

    if (pos < mid)

    {

        modify(pos, value, idx*2, left, mid);

    }

    else

    {

        modify(pos, value, idx*2+1, mid, right);

    }

}

int main()

{

    a.pb(-1); // Ignora a posição 0

    // Elementos do array

    a.pb(1);

    a.pb(2);

    a.pb(3);

    a.pb(4);

    a.pb(5);

    a.pb(6);

    st.resize(a.size() *4); // 4n para o pior caso

    build();

    for (auto it : st)

    {

        cout << it <<" ";

    }

    cout << endl;

    cout << sum(1,5) << endl; // espera-se 1+2+3+4 = 10

    modify(2,10); // passou de 2 para 10

    cout << sum(1,5) << endl; // espera-se: 1+10+3+4 = 18

    return 0;

}

Range Minimal Query (RMQ) usando uma Segment Tree