Árbol Wavelet

Introducción

Supongamos que tenemos un arreglo estático de enteros $a_0, a_1, \dots, a_{N-1}$ que cumple $0\le a_i<\sigma$, y queremos responder consultas online de la siguiente forma:

• Hallar el $k$-ésimo elemento más pequeño en el subarreglo contiguo $a[l:r)$ (donde $k$ está indexado desde 0).

En este módulo, introduciremos el concepto de un Árbol Wavelet (Wavelet Tree) para responder estas consultas de forma eficiente tanto en tiempo como en memoria. Cada solución del módulo se construye sobre la anterior.

Recursos

Leer estos recursos es opcional, a menos que las explicaciones del módulo resulten demasiado sucintas.

K-ésimo más pequeño en un rango: Solución 1

Empecemos construyendo un Árbol de Segmentos sobre los valores $[0, \sigma)$. Un nodo del árbol correspondiente a un rango de valores $[v_l, v_r)$ guardará

1. Una lista que contiene los índices del arreglo $a$ con valores en ese rango, en orden creciente.
2. Si el nodo no es una hoja (es decir, $v_l + 1 < v_r$), punteros a sus dos nodos hijos, correspondientes a los rangos $[v_l, (v_l+v_r)/2)$ y $[(v_l+v_r)/2, v_r)$.

Un árbol donde cada nodo guarda una lista de todo lo que tiene debajo en orden se llama un Merge Sort Tree.

Para construir esta estructura de datos, empezamos en el nodo raíz correspondiente al rango $[0,\sigma)$, partimos los índices $[0,N)$ entre sus dos hijos y construimos recursivamente cada hijo. Esto toma $O(N\log \sigma)$ en tiempo y memoria.

Para responder una consulta, de nuevo empezamos en el nodo raíz y luego caminamos recursivamente por el árbol hasta llegar al nodo hoja correspondiente al valor respuesta. Para determinar si caminar hacia el hijo izquierdo o el derecho del nodo actual, primero consultamos el número de índices en el vector de índices del hijo izquierdo en el rango $[l,r)$ y lo guardamos en una variable \texttt{num\\_left}.

• Si k<\texttt{num\\_left}, entonces la respuesta es el $k$-ésimo valor más pequeño en el hijo izquierdo.
• En caso contrario, la respuesta es el (k-\texttt{num\\_left})-ésimo valor más pequeño en el hijo derecho.

Consultar el conteo en un solo nodo toma $O(\log N)$ usando búsqueda binaria, y el árbol tiene profundidad $O(\log \sigma)$, así que en total una consulta toma $O(\log \sigma\log N)$.

Implementación

Nota: fijamos $\sigma=2^{30}$ para que cada nodo tenga longitud igual a una potencia de dos.

#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;
	void build(const vector<int> &A, int b) {
		if (b == 0 || inds.empty()) return;
		l = new Wavelet();
		r = new Wavelet();
		for (int x : inds) {
			if (A[x] & (1 << (b - 1))) r->inds.push_back(x);
			else l->inds.push_back(x);
		}
		l->build(A, b - 1);
		r->build(A, b - 1);
	}
	// k-th (0-indexed) smallest value, only considering A[l, r)
	int range_kth_smallest(int l, int r, int k, int b) {
		if (b == 0) return 0;
		int num_left = count_prefix(this->l->inds, r) - count_prefix(this->l->inds, l);
		if (k < num_left) return this->l->range_kth_smallest(l, r, k, b - 1);
		return (1 << (b - 1)) + this->r->range_kth_smallest(l, r, k - num_left, b - 1);
	}
};
 
int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
 
	int N, Q;
	cin >> N >> Q;
	vector<int> A(N);
	for (int &a : A) cin >> a;
 
	// build tree
	Wavelet *root = new Wavelet();
	root->inds.resize(N);
	iota(begin(root->inds), end(root->inds), 0);
	const int MAX_BIT = 30;
	root->build(A, MAX_BIT);
 
	// answer queries
	for (int q = 0; q < Q; ++q) {
		int l, r, k;
		cin >> l >> r >> k;
		cout << root->range_kth_smallest(l, r, k, MAX_BIT) << "\n";
	}
}

K-ésimo más pequeño en un rango: Solución 2a

Nuestro objetivo en esta sección es quitar el factor $\log N$ de la complejidad temporal de consulta de la solución 1 sin cambiar la complejidad espacial. Seguiremos guardando un vector de enteros en cada nodo del Árbol de Segmentos. Su longitud será la misma que antes, pero representará algo distinto.

Consideremos primero qué vector deberíamos guardar en el nodo raíz para calcular \texttt{num\\_left} sin búsqueda binaria. Lo más simple que podríamos hacer es guardar los valores de \texttt{count\\_prefix}(\texttt{this->l->inds}, r) para cada $r$ posible de $0$ a $N$ inclusive. Es decir, todas las sumas de prefijos del bitvector de longitud $N$ con el $i$-ésimo elemento igual a $1$ si $a_i$ mapea al nodo hijo izquierdo, y $0$ en caso contrario. Entonces \texttt{num\\_left} se puede calcular en tiempo constante simplemente restando dos sumas de prefijos.

En general, en cada nodo que no es hoja del Árbol de Segmentos, podemos primero construir un vector de bits de longitud igual a la subsecuencia de $A$ asociada con ese nodo, con $1$s para valores que mapean al hijo izquierdo y $0$s para los demás, y luego guardar sus sumas de prefijos en ese nodo.

Para responder consultas, a diferencia de la solución 1, necesitaremos modificar $l$ y $r$ a medida que caminamos por el árbol. En lugar de representar los índices $l$-ésimo a $r$-ésimo de $A$, ahora representarán los índices $l$-ésimo a $r$-ésimo de la subsecuencia de $A$ asociada con el nodo actual.

Implementación

Nota: la implementación evita guardar \texttt{count\\_prefix}(\texttt{this->l->inds}, 0) ya que siempre es cero.

#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;
		l = new Wavelet();
		r = new Wavelet();
		int num_left = 0;
		vector<int> A0, A1;
		for (int x : A) {
			if (x & (1 << (b - 1))) {
				A1.push_back(x);
			} else {
				++num_left;
				A0.push_back(x);
			}
			num_lefts.push_back(num_left);
		}
		l->build(A0, b - 1);
		r->build(A1, b - 1);
	}
	int range_kth_smallest(int l, int r, int k, int b) {
		if (b == 0) return 0;
		int pr = count_prefix(num_lefts, r);
		int pl = count_prefix(num_lefts, l);
		int num_left = pr - pl;
		if (k < num_left) return this->l->range_kth_smallest(pl, pr, k, b - 1);
		return (1 << (b - 1)) +
		       this->r->range_kth_smallest(l - pl, r - pr, k - num_left, b - 1);
	}
};
 
int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
 
	int N, Q;
	cin >> N >> Q;
	vector<int> A(N);
	for (int &a : A) cin >> a;
 
	// build tree
	Wavelet *root = new Wavelet();
	const int MAX_BIT = 30;
	root->build(A, MAX_BIT);
 
	// answer queries
	for (int q = 0; q < Q; ++q) {
		int l, r, k;
		cin >> l >> r >> k;
		cout << root->range_kth_smallest(l, r, k, MAX_BIT) << "\n";
	}
}

Solución 2b

La siguiente solución tiene la misma complejidad temporal y espacial que la anterior, pero el factor constante es mucho mejor. Específicamente, es más del doble de rápida y usa menos de una décima parte de la memoria.

Para lograr esto, concatena los vectores de bits de cada nivel en un solo vector de bits de longitud $N$ antes de tomar sumas de prefijos. Esta construcción se conoce como la Matriz Wavelet (Wavelet Matrix).

Se puede ver que el proceso de consulta es equivalente al de la solución de arriba (salvo traducir $l$ y $r$ por una constante).

#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);
	for (int &a : A) cin >> a;
 
	const int MAX_BIT = 30;
	vector<vector<int>> num_lefts(MAX_BIT);
	for (int b = MAX_BIT - 1; b >= 0; --b) {
		vector<int> A0, A1;
		int num_left = 0;
		for (int x : A) {
			if (x & (1 << b)) {
				A1.push_back(x);
			} else {
				++num_left;
				A0.push_back(x);
			}
			num_lefts.at(b).push_back(num_left);
		}
		swap(A, A0);
		A.insert(end(A), begin(A1), end(A1));
	}
 
	auto get_num_left = [&](int b, int r) {
		return r == 0 ? 0 : num_lefts.at(b).at(r - 1);
	};
	auto range_kth_smallest = [&](int l, int r, int k) {
		int ans = 0;
		for (int b = MAX_BIT - 1; b >= 0; --b) {
			int pr = get_num_left(b, r);
			int pl = get_num_left(b, l);
			int num_left = pr - pl;
			if (k < num_left) {
				l = pl, r = pr;
			} else {
				k -= num_left;
				l = l - pl + num_lefts.at(b).back();
				r = r - pr + num_lefts.at(b).back();
				ans ^= 1 << b;
			}
		}
		return ans;
	};
	for (int q = 0; q < Q; ++q) {
		int l, r, k;
		cin >> l >> r >> k;
		cout << range_kth_smallest(l, r, k) << "\n";
	}
}

K-ésimo más pequeño en un rango: Solución 3

Aquí discutimos cómo quitar el factor de $\log \sigma$ de la complejidad espacial.

El cuello de botella de memoria en la solución 2 es guardar las sumas de prefijos de $\log \sigma$ vectores de bits de longitud $N$, lo que toma $O(N\log \sigma)$ enteros con el enfoque más directo. Sin embargo, si podemos reducir esto a $O(N\log \sigma)$ bits de información, podemos empaquetar estos bits en $O(N\log \sigma / W)$ palabras donde $W$ es el tamaño de palabra  ($W=64$ en una arquitectura de 64 bits). Esto son $O(N)$ palabras asumiendo $2^W>\sigma$ (es decir, todos los enteros con los que trabajamos caben en una sola palabra).

Resta describir cómo guardar un solo vector de bits de longitud $N$ en $O(N)$ bits permitiendo todavía consultas de suma de prefijos en tiempo constante. Específicamente, podemos guardar el vector de bits original en $N$ bits y solo las sumas de prefijos con longitud divisible por $W$, tomando $O(N\log N/W)$ bits, que son $O(N)$ bits asumiendo $2^W>N$. Para responder una consulta de la $r$-ésima suma de prefijos en tiempo constante, empezamos con la $\lfloor r/W\rfloor\cdot W$-ésima suma de prefijos y luego usamos operaciones built-in que corren en tiempo constante para sumar la contribución de los $r\%W$ bits restantes (como \texttt{\\_\\_builtin\\_popcountll} para contar el número de bits activados en una palabra de 64 bits).

Implementación

#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) {
			if (v.at(i)) packed.at(i / BITS) |= 1ULL << (i % BITS);
		}
		psums = {0};
		for (auto b : packed) psums.push_back(psums.back() + __builtin_popcountll(b));
	}
	int count_prefix(int r) {
		return psums.at(r / BITS) +
		       __builtin_popcountll(packed.at(r / BITS) & ((1ULL << (r % BITS)) - 1));
	}
	int count() { return psums.back(); }
};
 
int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
 
	int N, Q;
	cin >> N >> Q;
	vector<int> A(N);
	for (int &a : A) cin >> a;
 
	const int MAX_BIT = 30;
	vector<PrefixSummer> num_lefts(MAX_BIT);
	for (int b = MAX_BIT - 1; b >= 0; --b) {
		vector<int> A0, A1;
		vector<bool> bitvec;
		for (int x : A) {
			if (x & (1 << b)) {
				bitvec.push_back(0);
				A1.push_back(x);
			} else {
				bitvec.push_back(1);
				A0.push_back(x);
			}
		}
		num_lefts.at(b).init(bitvec);
		swap(A, A0);
		A.insert(end(A), begin(A1), end(A1));
	}
 
	auto range_kth_smallest = [&](int l, int r, int k) {
		int ans = 0;
		for (int b = MAX_BIT - 1; b >= 0; --b) {
			int pr = num_lefts.at(b).count_prefix(r);
			int pl = num_lefts.at(b).count_prefix(l);
			int num_left = pr - pl;
			if (k < num_left) {
				l = pl, r = pr;
			} else {
				k -= num_left;
				l = l - pl + num_lefts.at(b).count();
				r = r - pr + num_lefts.at(b).count();
				ans ^= 1 << b;
			}
		}
		return ans;
	};
	for (int q = 0; q < Q; ++q) {
		int l, r, k;
		cin >> l >> r >> k;
		cout << range_kth_smallest(l, r, k) << "\n";
	}
}

Problemas