Tutorial

Wavelet Matrix

Artigo explicando a estrutura de wavelet matrix e o que pode ser feito com ela.

Requisitos

• Trie binária

Introdução

Com uma trie binária, é possível suportar os seguintes tipos de operação:

• insert(x), insere um elemento $x$ na estrutura. $O(H)$
• countLess(a), calcular a quantidade dos elementos menores que $a$. $O(H)$.
• sumLess(a) calcular a soma dos elementos menores que $a$. $O(H)$.
• maxXor(a) calcular o maior valor de $a\ \mathrm{XOR}\ x$ onde $x$ é um elemento na estrutura de dados. $O(H)$.

E uma trie binária pode funcionar efetivamente como uma Segment Tree dinâmica.

Uma Wavelet Matrix é como uma vesão diferente de uma trie binária que pode suportar os seguintes tipos de operação em um array estático:

• rangeKth(l, r, k), encontrar o $k$-ésimo menor elemento no range $[l, r)$. $O(H)$.
• rangeCntLess(l, r, x), calcular quantos elementos no intervalo $[l, r)$ são menores que $x$. $O(H)$.
• rangeKthXor(l, r, k, x), encontrar o $k$-ésimo menor valor de $a_i\ \mathrm{XOR}\ x$ onde $l \le i < r$. $O(H)$.

A estrutura

Considere um array $A = (a_0, a_1, \dots, a_{n - 1})$. Assim como em uma trie binária, vamos olhar todos os bits dos números, do mais significante até o menos significante.

Digamos que a matriz tem altura $H$. Cada elemento do array deve estar no intervalo $[0, 2^H - 1]$. Para a primeira linha da matriz, vamos colocar o $H$-ésimo bit de cada elemento (do menos significante ao mais significante).

Exemplo: $A = (2, 6, 1, 5, 7, 4), H = 3$. A primeira linha da matriz é $(0, 1, 0, 1, 1, 1)$.

A partir da segunda linha, os elementos que foram 0 na última linha são movidos para a esquerda, e os que eram 1 na última linha são movidos para a direita. No exemplo acima, a matriz é reordenada na seguinte forma: $(2, 1 | 6, 5, 7, 4)$.

O processo continua como na primeira linha, com os próximos bits do número.

A matriz final no exemplo acima é:

É evidente pela estrutura que o array original pode ser recuperado a partir da matriz. Para realizar operações de forma eficiente, nós mantemos uma matriz de soma de prefixos, em vez da matriz de 0s e 1s. Com isso, podemos calcular a quantidade de 0s e 1s em intervalos arbitrários em $O(1)$.

Implementação

Segue uma simples implementação da estrutura para propósitos educacionais.

struct WaveletMatrix {
    int n, h;
    vector<vector<int>> mat;
    WaveletMatrix(vector<int> &a, int H) : n(a.size()), h(H), mat(H, vector<int>(n + 1)) {
        vector<int> row = a;
        for (int i = h - 1; i >= 0; i--) {
            vector<int> left, right;
            for (int j = 1; j <= n; j++) {
                mat[i][j] = row[j - 1] >> i & 1;
                if (mat[i][j]) right.push_back(row[j - 1]);
                else left.push_back(row[j - 1]);
            }
            row.clear();
            for (int x: left) row.push_back(x);
            for (int x: right) row.push_back(x);
            for (int j = 1; j <= n; j++) {
                mat[i][j] += mat[i][j - 1];
            }
        }
    }

    tuple<int, int, int, int> split(int i, int l, int r) {
        int a1 = mat[i][l], a0 = l - a1;
        int b1 = mat[i][r], b0 = r - b1;
        int zeros = n - mat[i][n];
        return {a0, b0, zeros + a1, zeros + b1};
    }

    int rangeKth(int i, int l, int r, int k) {
        if (i == 0) return 0;
        auto [l0, r0, l1, r1] = split(i - 1, l, r);
        int cntZero = r0 - l0;
        if (cntZero < k) {
            return rangeKth(i - 1, l0, r0, k);
        } else {
            return (1 << (i - 1)) + rangeKth(i - 1, l1, r1, k - cntZero);
        }
    }

    int rangeCntLess(int i, int l, int r, int x) {
        if (i == 0) return 0;
        auto [l0, r0, l1, r1] = split(i - 1, l, r);
        int cntZero = r0 - l0;
        if (x >> (i - 1) & 1) {
            return cntZero + rangeCntLess(i - 1, l1, r1, x);
        } else {
            return rangeCntLess(i - 1, l0, r0, x);
        }
    }
};

Cada intervalo em uma linha é quebrado em dois intervalos na linha seguinte, uma com os 0s e uma com os 1s. A função split encontra os dois intervalos na linha seguinte que representam um intervalo na linha atual.

Na operação rangeKth, começamos com um intervalo no array original, e verificamos, para cada bit do mais significativo ao menos, se o $k$-ésimo menor valor do intervalo tem um 0 ou um 1 na posição atual. Os 0s sempre são menores que os 1s, então basta contar quantos 0s e quantos 1s no dígito binário atual.

A operação rangeCntLess funciona de forma similar.

Como exercício, sugerimos que o leitor implemente a operação rangeKthXor, como descrita na introdução.

Problemas

• Range Kth Smallest