Skip to Content
CTRL K

Rectangle Sum

Explicación

Nos dan NN puntos en un plano 2D. Cada punto tiene coordenadas (xi,yi)(x_i, y_i) y un peso wiw_i. Para cada consulta (l, d, r, u), debemos hallar la suma de los pesos de todos los puntos dentro del rectángulo:

l ≤ x < r y d ≤ y < u.

Observación inicial

Una observación útil es que la restricción del rectángulo se parte en un rango de x y un rango de y. Si ordenamos todos los puntos por su coordenada x, la condición l≤x<rl ≤ x < r se convierte en un rango contiguo de índices en el arreglo ordenado.

Entonces el problema queda:

Entre los puntos en el rango de índices [L, R], hallar la suma de los pesos cuyo y está en [d, u).

Para soportar consultas rápidas sobre la coordenada y, comprimimos todos los valores de y para que queden en un rango pequeño [1 ... M].

Luego construimos un Árbol Wavelet sobre el arreglo de puntos (ordenados por x). El valor guardado en el árbol es la coordenada y comprimida, mientras que el dato asociado es el peso del punto.

¿Cómo modificamos la plantilla normal de Árbol Wavelet?

Normalmente, un Árbol Wavelet puede contar cuántos valores caen dentro de un rango.

Pero nuestro problema pide la suma de pesos, no solo el recuento. Para soportar esto, guardamos el arreglo adicional pref, donde:

pref[i] = suma de pesos de los primeros i elementos en este nodo

Por lo tanto, cada nodo guarda:

  • b — arreglo de prefijos que indica cuántos elementos del prefijo actual van al hijo izquierdo
  • pref — una suma de prefijos de pesos de los elementos de este segmento

El arreglo b nos permite traducir un rango de consulta [L, R] del nodo actual a los rangos correctos en sus hijos al descender por el árbol.

Procesar una consulta

Cuando el rango de y del nodo actual queda completamente dentro del rango de consulta [D, U], podemos devolver de forma directa:

pref[R] - pref[L - 1]

Esto evita bajar más en el árbol y da la suma de pesos de inmediato.

Durante una consulta:

  1. Convertimos el rectángulo al rango de índices [L, R] usando x.
  2. Convertimos los límites de y a índices comprimidos [D, U].
  3. Recorremos el Árbol Wavelet, usando b para mapear índices y pref para acumular sumas de pesos.

Como la altura del árbol es O(log N), cada consulta se procesa de forma eficiente.

Implementación

Complejidad temporal: O((N+Q)log⁡N)\mathcal{O}((N + Q)\log N)

 
#include <bits/stdc++.h>
using namespace std;
 
// BeginCodeSnip{Our Wavelet Tree template}
struct WaveletTree {
	int lo, hi;
	WaveletTree *l, *r;
	vector<int> b;
	vector<long long> pref;
 
	WaveletTree(pair<int, long long> *from, pair<int, long long> *to, int x, int y) {
		lo = x;
		hi = y;
		l = r = nullptr;
 
		pref.reserve(to - from + 1);
		pref.push_back(0);
		for (auto it = from; it != to; ++it) pref.push_back(pref.back() + it->second);
 
		if (lo == hi || from >= to) return;
 
		int mid = (lo + hi) >> 1;
		auto f = [&](const pair<int, long long> &e) { return e.first <= mid; };
 
		b.reserve(to - from + 1);
		b.push_back(0);
		for (auto it = from; it != to; ++it) b.push_back(b.back() + f(*it));
 
		auto p = stable_partition(from, to, f);
		l = new WaveletTree(from, p, lo, mid);
		r = new WaveletTree(p, to, mid + 1, hi);
	}
 
	long long query(int L, int R, int d, int u) {
		if (L > R || u < lo || d > hi) return 0;
		if (d <= lo && hi <= u) return pref[R] - pref[L - 1];
 
		int lb = b[L - 1], rb = b[R];
		long long res = 0;
 
		if (l) res += l->query(lb + 1, rb, d, u);
		if (r) res += r->query(L - lb, R - rb, d, u);
 
		return res;
	}
};
 
struct Point {
	int x, y;
	long long w;
};
// EndCodeSnip
 
int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
 
	// BeginCodeSnip{Getting input and Building the Wavelet tree}
	int n, q;
	cin >> n >> q;
 
	vector<Point> p(n);
	vector<int> ys(n);
 
	for (int i = 0; i < n; i++) {
		cin >> p[i].x >> p[i].y >> p[i].w;
		ys[i] = p[i].y;
	}
 
	sort(ys.begin(), ys.end());
	ys.erase(unique(ys.begin(), ys.end()), ys.end());
 
	sort(p.begin(), p.end(), [](const Point &a, const Point &b) { return a.x < b.x; });
 
	vector<pair<int, long long>> a(n);
	vector<int> xs(n);
 
	for (int i = 0; i < n; i++) {
		int y = lower_bound(ys.begin(), ys.end(), p[i].y) - ys.begin() + 1;
		a[i] = {y, p[i].w};
		xs[i] = p[i].x;
	}
 
	WaveletTree wt(a.data(), a.data() + n, 1, ys.size());
	// EndCodeSnip
 
	int l, d, r, u, L, D, R, U;
	while (q--) {
 
		cin >> l >> d >> r >> u;
 
		L = lower_bound(xs.begin(), xs.end(), l) - xs.begin() + 1;
		R = lower_bound(xs.begin(), xs.end(), r) - xs.begin();
		D = lower_bound(ys.begin(), ys.end(), d) - ys.begin() + 1;
		U = lower_bound(ys.begin(), ys.end(), u) - ys.begin();
 
		if (L > R || D > U) cout << 0 << '\n';
		else cout << wt.query(L, R, D, U) << '\n';
	}
}
Fundamentos