Wavelet Tree

Introduction

Suppose we have a static array of integers $a_0, a_1, \dots, a_{N-1}$ satisfying $0\le a_i<\sigma$, and we want to answer online queries of the following form:

• Find the $k$-th smallest element in the contiguous subarray $a[l:r)$ (where $k$ is 0-indexed).

In this module, we will introduce the concept of a Wavelet Tree to answer these queries efficiently with respect to both time and memory. Each module solution will build on the previous one.

Resources

Reading these resources is optional, unless you find the in-module explanations too succinct.

Range Kth Smallest: Solution 1

Let's start by building a segment tree on the values $[0, \sigma)$. A segment tree node corresponding to a range of values $[v_l, v_r)$ will store

1. A list containing the indices of the array $a$ with values in that range, in increasing order.
2. If the node is not a leaf (that is, $v_l + 1 < v_r$), pointers to its two child nodes, corresponding to the ranges $[v_l, (v_l+v_r)/2)$ and $[(v_l+v_r)/2, v_r)$.

A tree where each node stores a list of everything under it in sorted order is called a merge-sort tree.

To build this data structure, we start at the root node corresponding to the range $[0,\sigma)$, partition the indices $[0,N)$ among its two children, and recursively build each child. This takes $O(N\log \sigma)$ time and memory.

To answer a query, we again start at the root node, then recursively walk down the tree until we reach the leaf node corresponding to the answer value. To determine whether to walk down into the left child or the right child of the current node, we first query the number of indices in the index vector of the left child in the range $[l,r)$ and store it into a variable $\texttt{num\_left}$.

• If $k<\texttt{num\_left}$, then the answer is the $k$th-smallest value in the left child.
• Otherwise, the answer is the $(k-\texttt{num\_left})$-th smallest value in the right child.

Querying the count in a single node takes $O(\log N)$ time using binary search, and the tree has depth $O(\log \sigma)$, so in a total a query takes $O(\log \sigma\log N)$ time.

Implementation

Note: we set $\sigma=2^{30}$ so that every node has length equal to a power of two.

Copy
#include <bits/stdc++.h>
using namespace std;

int count_prefix(const vector<int> &v, int r) {
	return lower_bound(begin(v), end(v), r) - begin(v);
}

struct Wavelet {
	vector<int> inds;
	Wavelet *l, *r;

Range Kth Smallest: Solution 2a

Our goal in this section is to remove the $\log N$ factor from the query time complexity of solution 1 without changing the memory complexity. We'll continue to store a vector of integers at each segment tree node. Its length will be the same as before, but it will represent something different.

Let's first consider what vector we should store at the root node to compute $\texttt{num\_left}$ without binary search. The simplest thing we could do is to store the values of $\texttt{count\_prefix}(\texttt{this->l->inds}, r)$ for each possible $r$ from $0$ to $N$ inclusive. That is, all prefix sums of the length-$N$ bitvector with $i$th element equal to $1$ if $a_i$ maps to the left child node, and $0$ otherwise. Then $\texttt{num\_left}$ can be computed in constant time just by subtracting two prefix sums.

In general, at each non-leaf node of the segment tree, we can first construct a bit vector of length equal to the subsequence of $A$ associated with that node with $1$s for values that map to the left child node and $0$s for others, and then store its prefix sums at that node.

To answer queries, unlike solution 1, we'll need to modify $l$ and $r$ as we walk down the tree. Instead of representing the $l$th through $r$th indices of $A$, they'll now represent the $l$th through $r$th indices of the subsequence of $A$ associated with the current node.

Implementation

Note: The implementation avoids storing $\texttt{count\_prefix}(\texttt{this->l->inds}, 0)$ since it's always zero.

Copy
#include <bits/stdc++.h>
using namespace std;

int count_prefix(const vector<int> &v, int r) { return r == 0 ? 0 : v.at(r - 1); }

struct Wavelet {
	vector<int> num_lefts;
	Wavelet *l, *r;
	void build(const vector<int> &A, int b) {
		if (b == 0 || A.empty()) return;

Solution 2b

The following solution has the same time and memory complexity as the previous one, but the constant factor is much better. Specifically, it's more than twice as fast, and uses less than one tenth the memory!

To accomplish this, it concatenates the bitvectors at each level into a single bitvector of length $N$ before taking prefix sums. This construction is known as the Wavelet Matrix.

The query process can be seen to be equivalent to that of the solution above (up to translating $l$ and $r$ by a constant).

Copy
#include <bits/stdc++.h>
using namespace std;

int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);

	int N, Q;
	cin >> N >> Q;
	vector<int> A(N);

Range Kth Smallest: Solution 3

Here we discuss how to remove the factor of $\log \sigma$ from the memory complexity.

The memory bottleneck in solution 2 is storing the prefix sums of $\log \sigma$ length-$N$ bitvectors, which takes $O(N\log \sigma)$ integers using the most straightforward approach. However, if we can reduce this to $O(N\log \sigma)$ bits of information, we can pack these bits into $O(N\log \sigma / W)$ words where $W$ is the word size ($W=64$ on a 64-bit architecture). This is $O(N)$ words assuming $2^W>\sigma$ (that is, all the integers we're working with fit into a single word).

It remains to describe how to store a single length-$N$ bitvector in $O(N)$ bits while still allowing constant time prefix sum queries. Specifically, we can store the original bitvector in $N$ bits and only the sums of prefixes with length divisible by $W$, taking $O(N\log N/W)$ bits, which is $O(N)$ bits assuming $2^W>N$. To answer a query for the $r$th prefix sum in constant time, we start with the $\lfloor r/W\rfloor\cdot W$th prefix sum and then use built-in operations that run in constant time to add the contribution of the remaining $r\%W$ bits (like $\texttt{\_\_builtin\_popcountll}$ to count the number of bits set in a 64-bit word).

Implementation

Copy
#include <bits/stdc++.h>
using namespace std;

struct PrefixSummer {
	const int BITS = 64;  // word size
	vector<uint64_t> packed;
	vector<int> psums;
	void init(const vector<bool> &v) {
		packed.resize(size(v) / BITS + 1);
		for (int i = 0; i < size(v); ++i) {

Problems