treap
(library/data-structure/treap.hpp)

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

使用方法

using S = long long;

S e() { return 0; }
S op(S a, S b) { return a + b; }
S reversal(S s) { return s; }

using F = tuple<>;

F id() { return {}; }
S mapping(F, S s) { return s; }
F composition(F, F) { return {}; }

using Treap = treap<S, e, op, reversal, F, id, mapping, composition>;

Treap tree, tree2, L, M, R;
Treap::iterator it;
S val;
int l, r, k;

tree.clear_tag();
tree.merge(tree2);

// std::multiset
int sz = tree.size();
bool empty = tree.empty();
tree.clear();
it = tree.insert(val);
it = tree.insert(it, val);
tree.erase(val);
tree.erase(it);
it = tree.lower_bound(val);
it = tree.upper_bound(val);
it = tree.find(val);

// sequence
int index = tree.get_index(it);
it = tree.insert_k(k, val);
tree.erase_k(k);
tie(L, R) = tree.split_k(k); // tree == empty
tie(L, M, R) = tree.split_range(l, r); // tree == empty
tree.reverse();
tree.reverse(l, r);

for(auto x : tree) {
	cout << x << endl; // call clear_tag() before iterating
}

題目

ABC306 E - Best Performances

CF Edu 15 F. T-Shirts

References

hitonanode’s library

tourist’s treap template

Depends on

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

Verified with

• test/data-structure/treap/unit-test-treap.test.cpp
• test/data-structure/treap/yosupo-Dynamic-Sequence-Range-Affine-Range-Sum.test.cpp
• test/data-structure/treap/yosupo-Range-Reverse-Range-Sum.test.cpp

Code

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


namespace felix {

namespace internal_treap {

template<class S, S (*e)(), S (*op)(S, S)>
struct treap_node {
	S key = e(), sum = e();
	int size = 1;
	treap_node* l = nullptr;
	treap_node* r = nullptr;
	treap_node* p = nullptr;

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

	void push() {}

	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;
		}
	}
};

} // namespace internal_treap


template<class S, S (*e)(), S (*op)(S, S), class Comp = std::less<S>>
struct treap : public internal::rbst_base<internal_treap::treap_node<S, e, op>, Comp> {
	using node_t = internal_treap::treap_node<S, e, op>;
	using base = internal::rbst_base<internal_treap::treap_node<S, e, op>, Comp>;
	using base::split_range_by_size, base::merge;
	using base::root;

public:
	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);
	}

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

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

} // namespace felix

#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 3 "library/data-structure/treap.hpp"

namespace felix {

namespace internal_treap {

template<class S, S (*e)(), S (*op)(S, S)>
struct treap_node {
	S key = e(), sum = e();
	int size = 1;
	treap_node* l = nullptr;
	treap_node* r = nullptr;
	treap_node* p = nullptr;

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

	void push() {}

	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;
		}
	}
};

} // namespace internal_treap


template<class S, S (*e)(), S (*op)(S, S), class Comp = std::less<S>>
struct treap : public internal::rbst_base<internal_treap::treap_node<S, e, op>, Comp> {
	using node_t = internal_treap::treap_node<S, e, op>;
	using base = internal::rbst_base<internal_treap::treap_node<S, e, op>, Comp>;
	using base::split_range_by_size, base::merge;
	using base::root;

public:
	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);
	}

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

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

} // namespace felix

Back to top page