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.

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.

Wavelet Tree Visualization

Let's say our array is: $[3,5,3,1,2,2,3,4,5,5]$ Each node has an array representing the indices of every number between $l$ and $r$

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$.

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 $\texttt{index}$[$i$] is smaller than mid otherwise $0$
• $\texttt{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 $\texttt{prefix\_b}$[$r$]
• If we recurse right, we use $r$ - $\texttt{prefix\_b}$[$r$]

Now let's try to solve our main problem.

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 that change the 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.

Let's first solve for adding a new element; erasing is similar but the opposite.

So what do the updates change?

• Our indices vector
• Our prefix vector

At each step of our recursion, the indices vector will need to be modified; we need to insert the new index. Since we can no longer maintain a sorted vector of the indices, we will switch to a set.

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, which is a binary search tree implementation in C++ that allows efficient queries for the rank of elements in a set. Querying for a prefix would then be equivalent to $\texttt{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