Proofs: Wavelet matrix

Access, rank, select and quantile on the wavelet matrix are correct full

This proof shows: correctness.

Notation as on the page. Symbols are $L$-bit integers; $c_\ell$ is bit $\ell$ of $c$ counted from the most significant, and $\text{pre}_\ell(c) = \lfloor c / 2^{L - \ell} \rfloor$ is the integer formed by its first $\ell$ bits. An occurrence is a position of $S$; we follow each occurrence through the levels.

Fact 1 (the move). Let $i$ be a position of $S_\ell$. The occurrence at $S_\ell[i]$ is at position $\text{down}_\ell(i)$ in $S_{\ell+1}$, with $\text{down}_\ell(i) = \text{rank}_0(B_\ell, i)$ if $B_\ell[i] = 0$ and $z_\ell + \text{rank}_1(B_\ell, i)$ if $B_\ell[i] = 1$. Moreover, for $0 \le i \le n$, the occurrences of $S_\ell[0 \ldots i-1]$ with bit 0 are exactly the first $\text{rank}_0(B_\ell, i)$ positions of $S_{\ell+1}$, and those with bit 1 are the positions $z_\ell, \ldots, z_\ell + \text{rank}_1(B_\ell, i) - 1$.

Proof. $S_{\ell+1}$ lists the occurrences with bit 0 in their order in $S_\ell$, in positions $0, \ldots, z_\ell - 1$, then those with bit 1 in their order, in positions $z_\ell, \ldots, n-1$. An occurrence with bit 0 at position $i$ is preceded among the 0-occurrences by those at the positions below $i$ with bit 0, which are $\text{rank}_0(B_\ell, i)$ many; symmetrically for 1. The second part is the same count for a prefix. $\square$

For $b \in \{0, 1\}$ write $\text{rank}'_b(i) = \text{rank}_0(B_\ell, i)$ if $b = 0$ and $z_\ell + \text{rank}_1(B_\ell, i)$ if $b = 1$; by Fact 1, $\text{down}_\ell(i) = \text{rank}'_{B_\ell[i]}(i)$.

Fact 2 (blocks). For each $\ell$ and each $\ell$-bit prefix $x$, the occurrences whose symbol has $\text{pre}_\ell = x$ form a contiguous block $[s_\ell(x), e_\ell(x))$ of $S_\ell$, listed in their order in $S$. Moreover, for $s_\ell(x) \le i \le e_\ell(x)$ and $b \in \{0, 1\}$, the occurrences of the block at positions below $i$ whose bit $\ell$ is $b$ are exactly those at the positions $[\text{rank}'_b(s_\ell(x)), \text{rank}'_b(i))$ of $S_{\ell+1}$; in particular $s_{\ell+1}(2x + b) = \text{rank}'_b(s_\ell(x))$.

Proof. By induction on $\ell$. For $\ell = 0$ the only prefix is empty and the block is all of $S_0 = S$. Assume the first claim for $\ell$ and take the block $[s, e)$ of $x$. The positions in $[s, i)$ with bit $b$ are, in increasing order, $i_1 < i_2 < \cdots < i_t$. By Fact 1, the positions of $S_\ell$ with bit $b$, taken in increasing order, go to consecutive positions of $S_{\ell+1}$ starting at $\text{rank}'_b(0)$, and those below a position $y$ go exactly to the positions below $\text{rank}'_b(y)$. So $\text{down}_\ell(i_1), \ldots, \text{down}_\ell(i_t)$ are the positions $\text{rank}'_b(s), \ldots, \text{rank}'_b(i) - 1$, in this order. With $i = e$, this shows that the occurrences with prefix $2x + b$, which are those of the block of $x$ with bit $b$, form a contiguous block of $S_{\ell+1}$ starting at $\text{rank}'_b(s)$, still in the order of $S$. This proves the first claim for $\ell + 1$ and the second claim for $\ell$. $\square$

Access. By Fact 1, lines L3 and L4 move $i$ to the position of the same occurrence one level down, and append its bit $\ell$ to $c$. After $L$ levels, $c$ has all $L$ bits of the symbol, from the most significant.

Rank. Let $i_0$ be the input, and $x_\ell = \text{pre}_\ell(c)$. Invariant before level $\ell$: $p = s_\ell(x_\ell)$, $p \le i \le e_\ell(x_\ell)$, and the occurrences of the block of $x_\ell$ at positions $[p, i)$ of $S_\ell$ are exactly the occurrences in $S[0 \ldots i_0 - 1]$ whose symbol has prefix $x_\ell$. At $\ell = 0$, $p = 0$, $i = i_0$ and the block is all of $S$. Lines L8 and L9 map $p$ and $i$ with $\text{rank}'_{c_\ell}$; by Fact 2, the new $[p, i)$ holds exactly those occurrences of $[p, i)$ whose bit $\ell$ is $c_\ell$, and $p = s_{\ell+1}(x_{\ell+1})$. After level $L - 1$, $x_L = c$, so $[p, i)$ holds exactly the occurrences of $c$ in $S[0 \ldots i_0 - 1]$, and line L10 returns their number.

Select. By the rank argument with $i_0 = 0$, line L11 gives $p = s_L(c)$. By Fact 2 the block of $c$ in $S_L$ lists the occurrences of $c$ in the order of $S$, so the $j$-th is at $q = p + j - 1$ (line L12); it is in the block because $j \le n_c$. Going up, if $c_\ell = 0$ the occurrence at position $q$ of $S_{\ell+1}$ came from a position $i$ of $S_\ell$ with $B_\ell[i] = 0$ and $\text{rank}_0(B_\ell, i) = q$, which is the $(q+1)$-th 0 of $B_\ell$ (line L14). If $c_\ell = 1$, it came from $i$ with $B_\ell[i] = 1$ and $z_\ell + \text{rank}_1(B_\ell, i) = q$, the $(q - z_\ell + 1)$-th 1 (line L15). After level 0, $q$ is the position in $S_0 = S$.

Quantile. Invariant before level $\ell$: $c = x$ is a prefix of $\ell$ bits, $[l, r)$ is a range inside the block of $x$ in $S_\ell$, the multiset of the symbols in $[l, r)$ is the multiset of the symbols of $S[l_0 \ldots r_0 - 1]$ with prefix $x$, and the answer is the $k$-th smallest of it. Initially $x$ is empty and the block is $S$. Among the symbols with prefix $x$, those with bit $\ell$ equal to 0 are all smaller than those with bit 1, and there are $t = \text{rank}_0(B_\ell, r) - \text{rank}_0(B_\ell, l)$ of them in $[l, r)$. If $k \le t$ the answer has bit 0 and is the $k$-th smallest of them; line L20 maps the range by Fact 2. Otherwise it has bit 1 and is the $(k - t)$-th smallest of the others, line L21. After $L$ levels, $c$ is the answer. $\blacksquare$