library/data-structures/1d-range-queries/associative-range-query.hpp

• View this file on GitHub
• Last update: 2022-07-21 16:12:33-04:00
• Include: #include "library/data-structures/1d-range-queries/associative-range-query.hpp"

Verified with

• verify/yosupo/yosupo-static_range_sum.test.cpp

Code

#pragma once

template <class T> struct RangeQuery {
	int sz; 
	std::vector<std::vector<T>> st;
	T id = 0;
	std::vector<T> a;

	T comb(T a, T b) { return a + b; }  

	void fill(int l, int r, int ind) {
		if (ind < 0) return;
		int m = (l + r) >> 1;
		T prod = id;
		for (int i = m - 1; i >= l; i--) 
			st[i][ind] = prod = comb(a[i], prod);
		prod = id;
		for (int i = m; i < r; i++) 
			st[i][ind] = prod = comb(prod, a[i]);
		fill(l, m, ind - 1);
		fill(m, r, ind - 1);
	}

	void init(std::vector<T> a_) {
		a = a_;
		sz = 1;
		while ((1 << sz) < (int)a.size()) sz++;
		st = std::vector<std::vector<T>>((1 << sz), std::vector<T>(sz));
		a.resize(1 << sz);
		fill(0, (1 << sz), sz - 1);
	}

	T query(int l, int r) {
		if (l == r) return a[l];
		int t = 31 - __builtin_clz(r ^ l);
		return comb(st[l][t], st[r][t]);
	}
};

template <class T> struct RangeQuery {
	int sz; 
	std::vector<std::vector<T>> st;
	T id = 0;
	std::vector<T> a;

	T comb(T a, T b) { return a + b; }  

	void fill(int l, int r, int ind) {
		if (ind < 0) return;
		int m = (l + r) >> 1;
		T prod = id;
		for (int i = m - 1; i >= l; i--) 
			st[i][ind] = prod = comb(a[i], prod);
		prod = id;
		for (int i = m; i < r; i++) 
			st[i][ind] = prod = comb(prod, a[i]);
		fill(l, m, ind - 1);
		fill(m, r, ind - 1);
	}

	void init(std::vector<T> a_) {
		a = a_;
		sz = 1;
		while ((1 << sz) < (int)a.size()) sz++;
		st = std::vector<std::vector<T>>((1 << sz), std::vector<T>(sz));
		a.resize(1 << sz);
		fill(0, (1 << sz), sz - 1);
	}

	T query(int l, int r) {
		if (l == r) return a[l];
		int t = 31 - __builtin_clz(r ^ l);
		return comb(st[l][t], st[r][t]);
	}
};

Back to top page