cpp_tree.snippets

snippet clssegtree "Segment tree class"
template <typename T, typename node_t = T>
class SegTree {
	const node_t defaultValue_;

 protected:
	virtual node_t transform(const T& v) const = 0;

	virtual node_t aggregate(const node_t& lhs, const node_t& rhs) const = 0;

 public:
	size_t n_;
	size_t width_;
	vector<node_t> tree_;

	virtual ~SegTree(){};

	SegTree(size_t n, const node_t& defaultValue)
			: defaultValue_(defaultValue),
				n_(n),
				width_(1 << bit_width(n - 1)),
				tree_(width_ * 2, defaultValue) {}

	void set(size_t i, const T& val) {
		tree_[i | width_] = transform(val);
	}

	void init() {
		for (size_t i = width_ - 1; i; --i) 
			tree_[i] = aggregate(tree_[i << 1], tree_[i << 1 | 1]);
	}

	void update(size_t i, const T& val) {
		tree_[i |= width_] = transform(val);
		while (i >>= 1)
			tree_[i] = aggregate(tree_[i << 1], tree_[i << 1 | 1]);
	}

	node_t query(size_t i) const { return tree_[i | width_]; }

	node_t query(size_t l, size_t r) const {
		node_t L = defaultValue_;
		node_t R = defaultValue_;
		for (l |= width_, r |= width_; l <= r; l >>= 1, r >>= 1) {
			if (l & 1)
				L = aggregate(L, tree_[l++]);
			if (~r & 1)
				R = aggregate(tree_[r--], R);
		}
		return aggregate(L, R);
	}
};

class RangeSumQuery : public SegTree<LL> {
	using SegTree<LL>::SegTree;

	LL transform(const LL& i) const { return i; }

	LL aggregate(const LL& l, const LL& r) const { return l + r; }
};
endsnippet

snippet clslazysegtree "Segment tree with lazy propagation"
template <typename T, typename node_t = T, typename lazy_t = T>
class LazySegTree {
	const node_t defaultValue_;
	const lazy_t identityValue_;

 protected:
	virtual node_t transform(const T& v) const = 0;

	virtual node_t aggregate(const node_t& lhs, const node_t& rhs) const = 0;

	virtual node_t applyLazy(const node_t& node, const lazy_t& diff) const = 0;

	virtual lazy_t compose(const lazy_t& lhs, const lazy_t& rhs) const = 0;

 public:
	size_t n_;
	size_t lg_;
	size_t width_;
	vector<node_t> tree_;
	vector<lazy_t> lazy_;

	LazySegTree(size_t n, const node_t& defaultValue, const lazy_t& identityValue)
			: defaultValue_(defaultValue),
				identityValue_(identityValue),
				n_(n),
				lg_(bit_width(n - 1)),
				width_(1 << lg_),
				tree_(width_ * 2, defaultValue),
				lazy_(width_, identityValue) {}

	virtual ~LazySegTree() {}

	void set(size_t i, const T& val) { tree_[i | width_] = transform(val); }

	void init() {
		for (size_t i = width_ - 1; i; --i)
			tree_[i] = aggregate(tree_[i << 1], tree_[i << 1 | 1]);
	}

	void update(size_t i, const lazy_t& f) {
		i |= width_;
		for (size_t j = lg_; j; j--)
			push(i >> j);
		apply(i, f);
		for (size_t j = 1; j <= lg_; j++)
			pull(i >> j);
	}

	void update(size_t l, size_t r, const lazy_t& f) {
		l |= width_, r |= width_;
		for (size_t i = lg_; i; --i) {
			if (l >> i << i != l)
				push(l >> i);
			if ((r + 1) >> i << i != r + 1)
				push(r >> i);
		}
		for (size_t L = l, R = r; L <= R; L >>= 1, R >>= 1) {
			if (L & 1)
				apply(L++, f);
			if (~R & 1)
				apply(R--, f);
		}
		for (size_t i = 1; i <= lg_; i++) {
			if (l >> i << i != l)
				pull(l >> i);
			if ((r + 1) >> i << i != r + 1)
				pull(r >> i);
		}
	}

	node_t query(size_t i) {
		i |= width_;
		for (int j = lg_; j; --j)
			push(i >> j);
		return tree_[i];
	}

	node_t query(size_t l, size_t r) {
		node_t L = defaultValue_, R = defaultValue_;
		l |= width_, r |= width_;
		for (size_t i = lg_; i; i--) {
			if (l >> i << i != l)
				push(l >> i);
			if ((r + 1) >> i << i != r + 1)
				push(r >> i);
		}
		for (; l <= r; l >>= 1, r >>= 1) {
			if (l & 1)
				L = aggregate(L, tree_[l++]);
			if (~r & 1)
				R = aggregate(tree_[r--], R);
		}
		return aggregate(L, R);
	}

 private:
	void apply(size_t i, const lazy_t& f) {
		tree_[i] = applyLazy(tree_[i], f);
		if (i < width_)
			lazy_[i] = compose(f, lazy_[i]);
	}

	void push(size_t i) {
		apply(i << 1, lazy_[i]);
		apply(i << 1 | 1, lazy_[i]);
		lazy_[i] = identityValue_;
	}

	void pull(size_t i) {
		tree_[i] = aggregate(tree_[i << 1], tree_[i << 1 | 1]);
	}
};

struct Node {
	LL val;
	size_t size;
};

class RangeSumQuery : public LazySegTree<LL, Node, LL> {
	using LazySegTree<LL, Node, LL>::LazySegTree;

	Node transform(const LL& i) const { return {i, 1}; }

	Node aggregate(const Node& l, const Node& r) const {
		return {l.val + r.val, l.size + r.size};
	}

	Node applyLazy(const Node& node, const LL& diff) const {
		return {node.val + diff * (LL)node.size, node.size};
	}

	LL compose(const LL& lhs, const LL& rhs) const { return lhs + rhs; }
};
endsnippet

snippet clsfenwicktree "Fenwick tree, binary indexed tree"
template <typename T>
class FenwickTree {
	const T defaultValue_;

 public:
	vector<T> arr_;
	vector<T> tree_;

	FenwickTree(int n, const T& defaultValue)
			: defaultValue_(defaultValue),
				arr_(n, defaultValue_),
				tree_(n, defaultValue_) {}

	void set(int i, const T& val) { arr_[i] = val; }

	void init() {
		tree_ = arr_;
		for (int i = 1, iEnd=(int)tree_.size(); i < iEnd; ++i) {
			int j = i + LSB(i);
			if (j < (int)tree_.size())
				tree_[j] += tree_[i];
		}
	}

	void update(int i, const T& val) {
		T diff = val - arr_[i];
		arr_[i] = val;
		while (i < tree_.size()) {
			tree_[i] += diff;
			i += LSB(i);
		}
	}

	T query(int i) {
		T res = defaultValue_;
		while (i > 0) {
			res += tree_[i];
			i -= LSB(i);
		}
		return res;
	}

	T query(int i, int j) {
		T res = defaultValue_;
		for (i--; j > i; j -= LSB(j))
			res += tree_[j];
		for (; i > j; i -= LSB(i))
			res -= tree_[i];
		return res;
	}

 private:
	static inline int LSB(int i) { return i & (-i); }
};
endsnippet