library/data-structure/lazy-treap.hpp

• View this file on GitHub
• Last update: 2023-07-11 10:41:16+08:00
• Include: #include "library/data-structure/lazy-treap.hpp"

Depends on

• library/bst/rbst-base.hpp
• library/random/rng.hpp

Code

#pragma once
#include <algorithm>

#include "../bst/rbst-base.hpp"


namespace felix {

namespace internal_treap {

template<class S,
         S (*e)(),
         S (*op)(S, S),
         S (*reversal)(S),
         class F,
         F (*id)(),
         S (*mapping)(F, S),
         F (*composition)(F, F)>
struct lazy_treap_node {
	S key, sum;
	F lz = id();
	int size = 1;
	bool rev = false;
	lazy_treap_node* l = nullptr;
	lazy_treap_node* r = nullptr;
	lazy_treap_node* p = nullptr;

	lazy_treap_node() : key(e()), sum(e()) {}
	lazy_treap_node(const S& s) : key(s), sum(s) {}

	void push() {
		if(lz != id()) {
			if(l != nullptr) {
				l->all_apply(lz);
			}
			if(r != nullptr) {
				r->all_apply(lz);
			}
			lz = id();
		}
		if(rev) {
			std::swap(l, r);
			if(l != nullptr) {
				l->rev ^= 1;
			}
			if(r != nullptr) {
				r->rev ^= 1;
			}
			sum = reversal(sum);
			rev = false;
		}
	}

	void pull() {
		size = 1;
		sum = key;
		if(l != nullptr) {
			size += l->size;
			sum = op(l->sum, sum);
			l->p = this;
		}
		if(r != nullptr) {
			size += r->size;
			sum = op(sum, r->sum);
			r->p = this;
		}
	}

	void all_apply(const F& f) {
		key = mapping(f, key);
		sum = mapping(f, sum);
		lz = composition(f, lz);
	}
};

} // namespace internal_treap


template<class S,
         S (*e)(),
         S (*op)(S, S),
         S (*reversal)(S),
         class F,
         F (*id)(),
         S (*mapping)(F, S),
         F (*composition)(F, F),
         class Comp = std::less<S>>
struct lazy_treap : public internal::rbst_base<internal_treap::lazy_treap_node<S, e, op, reversal, F, id, mapping, composition>, Comp> {
	using node_t = internal_treap::lazy_treap_node<S, e, op, reversal, F, id, mapping, composition>;
	using base = internal::rbst_base<node_t, Comp>;
	using base::split_range_by_size, base::merge;
	using base::root;

public:
	lazy_treap() {}

	void set(int k, const S& s) {
		auto [lhs, mid, rhs] = split_range_by_size(k, k + 1);
		mid.clear();
		mid.insert(mid.end(), s);
		merge(lhs), merge(mid), merge(rhs);
	}

	void apply(int l, int r, const F& f) {
		if(l == r) {
			return;
		}
		auto [lhs, mid, rhs] = split_range_by_size(l, r);
		mid.root->all_apply(f);
		merge(lhs), merge(mid), merge(rhs);
	}

	S prod(int l, int r) {
		if(l == r) {
			return e();
		}
		auto [lhs, mid, rhs] = split_range_by_size(l, r);
		mid.root->push();
		S ans = mid.root->sum;
		merge(lhs), merge(mid), merge(rhs);
		return ans;
	}

	S all_prod() {
		root->push();
		return root->sum;
	}

	void reverse(int l, int r) {
		if(l == r) {
			return;
		}
		auto [lhs, mid, rhs] = split_range_by_size(l, r);
		mid.root->rev ^= 1;
		merge(lhs), merge(mid), merge(rhs);
	}
};

} // namespace felix

#line 2 "library/data-structure/lazy-treap.hpp"
#include <algorithm>

#line 2 "library/bst/rbst-base.hpp"
#include <tuple>

#include <functional>

#line 2 "library/random/rng.hpp"
#include <chrono>

namespace felix {

inline unsigned long long rng() {
	static unsigned long long SEED = std::chrono::steady_clock::now().time_since_epoch().count();
	SEED ^= SEED << 7;
	SEED ^= SEED >> 9;
	return SEED;
}

} // namespace felix
#line 5 "library/bst/rbst-base.hpp"

namespace felix {

namespace internal {

template<class node_t, class Comp = std::less<decltype(node_t::key)>>
struct rbst_base {
	using key_type = decltype(node_t::key);

private:
	static const Comp Compare;

public:
	static node_t* merge(node_t* a, node_t* b) {
		if(a == nullptr || b == nullptr) {
			return a != nullptr ? a : b;
		}
		if((int) (((unsigned int) rng() * 1LL * (a->size + b->size)) >> 32) < a->size) {
			a->push();
			a->r = merge(a->r, b);
			a->pull();
			return a;
		} else {
			b->push();
			b->l = merge(a, b->l);
			b->pull();
			return b;
		}
	}

	static std::pair<node_t*, node_t*> split(node_t* v, int k) {
		if(v == nullptr) {
			return {nullptr, nullptr};
		}
		v->push();
		if(k <= get_size(v->l)) {
			auto p = split(v->l, k);
			v->l = p.second;
			if(p.first != nullptr) {
				p.first->p = nullptr;
			}
			v->pull();
			return {p.first, v};
		} else {
			auto p = split(v->r, k - get_size(v->l) - 1);
			v->r = p.first;
			if(p.second != nullptr) {
				p.second->p = nullptr;
			}
			v->pull();
			return {v, p.second};
		}
	}

	static std::tuple<node_t*, node_t*, node_t*> split_range(node_t* v, int l, int r) {
		auto p = split(v, l);
		auto q = split(p.second, r - l);
		return {p.first, q.first, q.second};
	}

	static void insert(node_t*& v, int k, const key_type& val) {
		auto p = split(v, k);
		v = merge(p.first, merge(new node_t(val), p.second));
	}

	static void erase(node_t*& v, int k) {
		auto p = split_range(v, k, k + 1);
		delete std::get<1>(p);
		v = merge(std::get<0>(p), std::get<2>(p));
	}

	static int lower_bound(node_t* v, const key_type& val) {
		if(v == nullptr) {
			return 0;
		}
		// TTTTFFFF

		//     ^

		if(Compare(v->key, val)) {
			return get_size(v->l) + 1 + lower_bound(v->r, val);
		} else {
			return lower_bound(v->l, val);
		}
	}

	static int upper_bound(node_t* v, const key_type& val) {
		// 1 2 3 3 4

		//         ^

		// F F F F T

		if(v == nullptr) {
			return 0;
		}
		if(!Compare(val, v->key)) {
			return get_size(v->l) + 1 + upper_bound(v->r, val);
		} else {
			return upper_bound(v->l, val);
		}
	}

	static key_type get(node_t*& v, int k) {
		auto p = split_range(v, k, k + 1);
		key_type res = std::get<1>(p)->key;
		v = merge(std::get<0>(p), merge(std::get<1>(p), std::get<2>(p)));
		return res;
	}

	static int get_index(node_t* v) {
		int k = get_size(v->l);
		while(v->p != nullptr) {
			if(v == v->p->r) {
				k++;
				if(v->p->l != nullptr) {
					k += v->p->l->size;
				}
			}
			v = v->p;
		}
		return k;
	}

	static void clear(node_t*& v) {
		postorder(v, [](node_t* v) {
			delete v;
		});
		v = nullptr;
	}

	static node_t* next(node_t* v) {
		if(v->r == nullptr) {
			while(v->p != nullptr && v->p->r == v) {
				v = v->p;
			}
			return v->p;
		}
		v->push();
		if(v->r == nullptr) {
			return nullptr;
		}
		v = v->r;
		while(v->l != nullptr) {
			v->push();
			v = v->l;
		}
		return v;
	}

	static node_t* prev(node_t* v) {
		if(v->l == nullptr) {
			while(v->p != nullptr && v->p->l == v) {
				v = v->p;
			}
			return v->p;
		}
		v->push();
		if(v->l == nullptr) {
			return nullptr;
		}
		v = v->l;
		while(v->r != nullptr) {
			v->push();
			v = v->r;
		}
		return v;
	}

	template<class Func>
	static void preorder(node_t* v, const Func& f) {
		auto rec = [&](auto rec, node_t* v) -> void {
			if(v == nullptr) {
				return;
			}
			v->push();
			f(v);
			rec(rec, v->l);
			rec(rec, v->r);
		};
		rec(rec, v);
	}

	template<class Func>
	static void inorder(node_t* v, const Func& f) {
		auto rec = [&](auto rec, node_t* v) -> void {
			if(v == nullptr) {
				return;
			}
			v->push();
			rec(rec, v->l);
			f(v);
			rec(rec, v->r);
		};
		rec(rec, v);
	}

	template<class Func>
	static void postorder(node_t* v, const Func& f) {
		auto rec = [&](auto rec, node_t* v) -> void {
			if(v == nullptr) {
				return;
			}
			v->push();
			rec(rec, v->l);
			rec(rec, v->r);
			f(v);
		};
		rec(rec, v);
	}

	static int size(node_t* v) { return get_size(v); }
	static bool empty(node_t* v) { return v == nullptr; }

private:
	static int get_size(node_t* v) { return v != nullptr ? v->size : 0; }
};

template<class node_t, class Comp> const Comp rbst_base<node_t, Comp>::Compare = Comp();

} // namespace internal


} // namespace felix

#line 4 "library/data-structure/lazy-treap.hpp"

namespace felix {

namespace internal_treap {

template<class S,
         S (*e)(),
         S (*op)(S, S),
         S (*reversal)(S),
         class F,
         F (*id)(),
         S (*mapping)(F, S),
         F (*composition)(F, F)>
struct lazy_treap_node {
	S key, sum;
	F lz = id();
	int size = 1;
	bool rev = false;
	lazy_treap_node* l = nullptr;
	lazy_treap_node* r = nullptr;
	lazy_treap_node* p = nullptr;

	lazy_treap_node() : key(e()), sum(e()) {}
	lazy_treap_node(const S& s) : key(s), sum(s) {}

	void push() {
		if(lz != id()) {
			if(l != nullptr) {
				l->all_apply(lz);
			}
			if(r != nullptr) {
				r->all_apply(lz);
			}
			lz = id();
		}
		if(rev) {
			std::swap(l, r);
			if(l != nullptr) {
				l->rev ^= 1;
			}
			if(r != nullptr) {
				r->rev ^= 1;
			}
			sum = reversal(sum);
			rev = false;
		}
	}

	void pull() {
		size = 1;
		sum = key;
		if(l != nullptr) {
			size += l->size;
			sum = op(l->sum, sum);
			l->p = this;
		}
		if(r != nullptr) {
			size += r->size;
			sum = op(sum, r->sum);
			r->p = this;
		}
	}

	void all_apply(const F& f) {
		key = mapping(f, key);
		sum = mapping(f, sum);
		lz = composition(f, lz);
	}
};

} // namespace internal_treap


template<class S,
         S (*e)(),
         S (*op)(S, S),
         S (*reversal)(S),
         class F,
         F (*id)(),
         S (*mapping)(F, S),
         F (*composition)(F, F),
         class Comp = std::less<S>>
struct lazy_treap : public internal::rbst_base<internal_treap::lazy_treap_node<S, e, op, reversal, F, id, mapping, composition>, Comp> {
	using node_t = internal_treap::lazy_treap_node<S, e, op, reversal, F, id, mapping, composition>;
	using base = internal::rbst_base<node_t, Comp>;
	using base::split_range_by_size, base::merge;
	using base::root;

public:
	lazy_treap() {}

	void set(int k, const S& s) {
		auto [lhs, mid, rhs] = split_range_by_size(k, k + 1);
		mid.clear();
		mid.insert(mid.end(), s);
		merge(lhs), merge(mid), merge(rhs);
	}

	void apply(int l, int r, const F& f) {
		if(l == r) {
			return;
		}
		auto [lhs, mid, rhs] = split_range_by_size(l, r);
		mid.root->all_apply(f);
		merge(lhs), merge(mid), merge(rhs);
	}

	S prod(int l, int r) {
		if(l == r) {
			return e();
		}
		auto [lhs, mid, rhs] = split_range_by_size(l, r);
		mid.root->push();
		S ans = mid.root->sum;
		merge(lhs), merge(mid), merge(rhs);
		return ans;
	}

	S all_prod() {
		root->push();
		return root->sum;
	}

	void reverse(int l, int r) {
		if(l == r) {
			return;
		}
		auto [lhs, mid, rhs] = split_range_by_size(l, r);
		mid.root->rev ^= 1;
		merge(lhs), merge(mid), merge(rhs);
	}
};

} // namespace felix

Back to top page