Proofs: Wavelet tree

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

This proof shows: correctness.

Notation as on the page: node $v$ covers $[a_v, b_v)$, $S_v$ is the subsequence of the symbols of $S$ in that range, in order, and for an inner node $m_v = \lfloor (a_v + b_v)/2 \rfloor$, $B_v[i] = [S_v[i] \ge m_v]$, with left child $\ell(v)$ covering $[a_v, m_v)$ and right child $r(v)$ covering $[m_v, b_v)$. Since $b_v - a_v \ge 2$ at an inner node, $a_v < m_v < b_v$, so both children cover nonempty ranges and the descent ends at a leaf after at most $\lceil \lg \sigma \rceil$ steps.

Lemma (mapping). Let $v$ be an inner node and $0 \le i \le |S_v|$. The symbols of $S_v[0 \ldots i-1]$ that are less than $m_v$, in order, are $S_{\ell(v)}[0 \ldots \text{rank}_0(B_v, i) - 1]$, and those that are at least $m_v$ are $S_{r(v)}[0 \ldots \text{rank}_1(B_v, i) - 1]$. In particular, if $i < |S_v|$ and $B_v[i] = 0$, then $S_v[i] = S_{\ell(v)}[\text{rank}_0(B_v, i)]$, and symmetrically for 1.

Proof. $S_{\ell(v)}$ is the subsequence of $S$ of the symbols in $[a_v, m_v)$, which is also the subsequence of $S_v$ of the symbols less than $m_v$, because $[a_v, m_v) \subseteq [a_v, b_v)$ and taking subsequences keeps order. The symbols of $S_v$ less than $m_v$ are those at the positions with $B_v = 0$. So the first $\text{rank}_0(B_v, i)$ of them are those at the positions with $B_v = 0$ below $i$. For the second claim, apply the first with $i + 1$: the symbol $S_v[i]$ is the last one of the prefix of length $\text{rank}_0(B_v, i+1) = \text{rank}_0(B_v, i) + 1$. The right side is symmetric. $\square$

Access. Invariant before each test of line L2: $v$ is the node covering $[a, b)$, $0 \le i < |S_v|$ and $S_v[i] = S[i_0]$, where $i_0$ is the input. It holds at the root. At an inner node, if $B_v[i] = 0$ then $S_v[i] < m$ and, by the lemma, $S_v[i] = S_{\ell(v)}[\text{rank}_0(B_v, i)]$; line L4 moves to exactly that node and position. Line L5 is symmetric. At a leaf, $[a, b) = \{a\}$, and every symbol of $S_v$ is $a$, so $S[i_0] = a$, which line L6 returns.

Rank. Let $c \in [0, \sigma)$ and $0 \le i_0 \le n$. Invariant: $c \in [a, b)$, $v$ covers $[a, b)$, and the number of $c$ in $S[0 \ldots i_0 - 1]$ equals the number of $c$ in $S_v[0 \ldots i - 1]$. At the root it is trivial. At an inner node with $c < m$, all copies of $c$ in $S_v[0 \ldots i-1]$ are less than $m$, so by the lemma they are exactly the copies of $c$ in $S_{\ell(v)}[0 \ldots \text{rank}_0(B_v, i) - 1]$; line L10 keeps the invariant, and line L11 is symmetric. At a leaf, every symbol of $S_v$ is $c$, so the count is $i$, which line L12 returns.

Select. Claim: for a node $v$ with $c \in [a_v, b_v)$ and $1 \le j \le$ the number of $c$ in $S_v$, SELECT$(v, a_v, b_v, c, j)$ returns the position in $S_v$ of the $j$-th $c$ of $S_v$. By induction on the height of $v$. At a leaf, all symbols are $c$ and the $j$-th is at $j - 1$ (line L13). At an inner node with $c < m_v$, the copies of $c$ in $S_v$ all lie in $S_{\ell(v)}$ in the same order, so by induction $p = $ SELECT$(\ell(v), \ldots)$ is the position in $S_{\ell(v)}$ of the $j$-th $c$. By the lemma, position $p$ of $S_{\ell(v)}$ is the position $q$ of $S_v$ with $B_v[q] = 0$ and $\text{rank}_0(B_v, q) = p$, which is the $(p+1)$-th 0 of $B_v$, that is $\text{select}_0(B_v, p + 1)$, line L15. Line L16 is symmetric. At the root, $S_v = S$.

Quantile. Invariant: $v$ covers $[a, b)$, the multiset $M$ of the symbols of $S_v[l \ldots r-1]$ is the multiset of the symbols of $S[l_0 \ldots r_0 - 1]$ that lie in $[a, b)$, and the answer, the $k_0$-th smallest of $S[l_0 \ldots r_0 - 1]$, is the $k$-th smallest of $M$ (in particular $1 \le k \le |M|$). Initially it holds. At an inner node, by the lemma applied to $r$ and to $l$, the symbols of $S_v[l \ldots r - 1]$ less than $m$ are those of $S_{\ell(v)}[\text{rank}_0(B_v, l) \ldots \text{rank}_0(B_v, r) - 1]$, and there are $z = \text{rank}_0(B_v, r) - \text{rank}_0(B_v, l)$ of them; the others are in the corresponding range of $S_{r(v)}$. In sorted order, the $z$ symbols less than $m$ come first. If $k \le z$, the $k$-th smallest of $M$ is the $k$-th smallest of the left part, line L20. Otherwise it is the $(k - z)$-th smallest of the right part, line L21. At a leaf $M$ consists of copies of $a$, and $k \le |M|$, so the answer is $a$, line L22. $\blacksquare$