Wavelet Tree

Wavelet trees are data structures that support efficient queries for the k-th minimum element in a range by maintaining a segment tree over values instead of indices.

Introduction

Suppose you want to support the following queries:

• Given a range, count the number of occurrences of value $x$.
• Given a range, find the $k$ smallest element

With a wavelet tree, you can easily support these queries in $\mathcal{O}(\log M)$ time, where $M$ is the maximum value in the array.

Wavelet Tree Structure

To answer value-based queries efficiently, we'll create a segment tree where each node represents a range of values, instead of indices. Just like a normal segment tree, each subsequent level splits the range into two halves. Note that an index can appear in at most $\log(M)$ nodes.

Solution - Range K-th Smallest

Before we solve this problem, let's consider a simpler version where we are asked, given a range, to count the number of occurrences of value $x$.

Solving the first type of query

Given a range $l$, $r$, count the number of occurrences of value x.

To calculate the number of occurrences from $𝑙$ to $𝑟$, we can use the following formula:

This reduces the problem to counting the number of occurrences in a prefix.

One way to solve the problem is to go to the leaf node and perform a binary search for the number of indices less than $𝑟$ However, let's explore a different approach that can also be extended to the second type of query.

Instead of binary searching on the leaf, we update $𝑟$ as we recurse down the tree. If we can determine the position (index) of $r$ in the left and right children of a node, we can recurse down the tree and determine its position in the leaf node.

To find the position of $𝑟$ in a node's left and right children, we need to determine how many indices are smaller than the middle value (mid) and precede $𝑟$. This can be done using a prefix sum.

Let's define:

• $c[i]$ = as $1$ if $index[i]$ is smaller than mid otherwise $0$
• $prefix_b[i]$ as prefix sum of $c[i]$

Formally

To update $r$ as we recurse down, we do the following:

• To know the value of $r$ if we recurse left, we use prefix_b[r]
• If we recurse right, we use $r$ - prefix_b[r]

Now let's try to solve our main problem.

Solving the second type of query

Given a range $l$, $r$ find the k-th smallest element

We will determine whether the answer for a given node is in the left or the right segment. We can calculate how many times the elements within the segments' ranges appear in our range $(l, r)$ using our first type of query. Note that this also works for non-leaf nodes using the following formula:

Therefore, the occurrences of the left node are

Note that $\texttt{left\_occurrences}$ is the number of indices between $l$ and $r$ whose value is less than $\texttt{mid}$

• If $\texttt{left\_occurrences}$ is greater or equal to $k$, it means the $k$-th smallest element is in the left subtree. Therefore, we update our range and recurse into the left child
• If $\texttt{left\_occurrences}$ is less than $k$, it means the $k$-th smallest element is in the right subtree. We adjust k by subtracting $\texttt{left\_occurrences}$ from $k$, update our range, and recurse into the right child

The answer then will be the value of the node we end up on (leaf).

In conclusion we maintain our ranges l and r as we recurse down to our child, and when we reach the child node we can return $r$ - $l$.

Implementation

Time Complexity: $\mathcal{O}(Q \cdot \log(M))$

Copy
#include <bits/stdc++.h>

using namespace std;
constexpr int MAX_VAL = 1e9 + 2;

struct Segment {
	Segment *left = nullptr, *right = nullptr;
	int l, r, mid;
	bool children = false;
	vector<pair<int, int>> indices;  // Index, Value

Supporting updates

Let's support updates of type:

• change value at index $i$ to $y$

We can traverse down to the leaf to remove the old element and also traverse down to add the new element.

To support erasing/adding values to the indices vector quickly, we can choose to use a set instead of a vector.

On the other hand, to change the prefix vector, since each update could change our prefix vector a lot, we can't maintain just the normal vector. What we could do is use a sparse segment tree.

• erasing and inserting can be done by just setting the value to 0 or 1 at the specific index
• querying for a prefix can be done by querying the segment tree from 0 to $i$

This approach is not memory efficient and requires a segment tree's implementation. A more friendly approach would be using an order statistics tree. Such that querying for a prefix would be equivalent to order_of_key($i$)

Implemention

Time Complexity: $\mathcal{O}(Q \cdot \log(M) \cdot \log(N))$

Copy
#include <bits/stdc++.h>

using namespace std;

#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
template <class T>
using Tree =
    tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

Problems