library/data-structure/range-add-range-sum.hpp

• View this file on GitHub
• Last update: 2023-07-24 08:41:31+08:00
• Include: #include "library/data-structure/range-add-range-sum.hpp"

Depends on

• library/data-structure/fenwick.hpp

Code

#pragma once
#include "fenwick.hpp"

namespace felix {

template<class T>
struct range_add_range_sum {
public:
	range_add_range_sum() {}
	explicit range_add_range_sum(int _n) : f0(_n), f1(_n) {}

	void add(int p, T x) { add(p, p + 1, x); }

	// [l, r)
	void add(int l, int r, T x) {
		f0.add(l, x);
		f1.add(l, -l * x);
		f0.add(r, -x);
		f1.add(r, r * x);
	}

	T get(int p) const { return sum(p, p + 1); }

	// [l, r)
	T sum(int l, int r) const {
		T res = 0;
		res += f0.get(r) * r + f1.get(r);
		res -= f0.get(l) * l + f1.get(l);
		return res;
	}

private:
	fenwick<T> f0, f1;
};

} // namespace felix

#line 2 "library/data-structure/fenwick.hpp"
#include <vector>
#include <cassert>

namespace felix {

template<class S>
struct fenwick {
public:
	fenwick() : n(0) {}
	explicit fenwick(int _n) : n(_n), data(_n) {}

	void add(int p, S x) {
		for(int i = p + 1; i <= n; i += i & -i) {
			data[i - 1] += x;
		}
	}

	// [0, p)
	S get(int p) const {
		auto ans = S();
		for(int i = p; i > 0; i -= i & -i) {
			ans += data[i - 1];
		}
		return ans;
	}

	// [l, r)
	S sum(int l, int r) const { return get(r) - get(l); }

	// 0-based
	int kth(S k) const {
		int x = 0;
		for(int i = 1 << std::__lg(n); i > 0; i >>= 1) {
			if (x + i <= n && k >= data[x + i - 1]) {
				x += i;
				k -= data[x - 1];
			}
		}
		return x;
	}

private:
	int n;
	std::vector<S> data;
};

} // namespace felix
#line 3 "library/data-structure/range-add-range-sum.hpp"

namespace felix {

template<class T>
struct range_add_range_sum {
public:
	range_add_range_sum() {}
	explicit range_add_range_sum(int _n) : f0(_n), f1(_n) {}

	void add(int p, T x) { add(p, p + 1, x); }

	// [l, r)
	void add(int l, int r, T x) {
		f0.add(l, x);
		f1.add(l, -l * x);
		f0.add(r, -x);
		f1.add(r, r * x);
	}

	T get(int p) const { return sum(p, p + 1); }

	// [l, r)
	T sum(int l, int r) const {
		T res = 0;
		res += f0.get(r) * r + f1.get(r);
		res -= f0.get(l) * l + f1.get(l);
		return res;
	}

private:
	fenwick<T> f0, f1;
};

} // namespace felix

Back to top page