pb_ds 平衡树 自定义 区间和node_update

• 如有错请指正

#include <iostream>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/detail/branch_policy/branch_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

#define TTT template<typename Node_CItr, typename Node_Itr, typename Cmp_Fn, typename _Alloc>

TTT class interval_sum_node_update : detail::branch_policy<Node_CItr, Node_Itr, _Alloc>
{
	private:
		typedef detail::branch_policy<Node_CItr, Node_Itr, _Alloc> base_type;
	public:
		typedef Cmp_Fn 									cmp_fn;
		typedef _Alloc 									allocator_type;
		typedef typename allocator_type::size_type 		size_type;
		typedef typename base_type::key_type 			key_type;
		typedef typename base_type::key_const_reference	key_const_reference;

		typedef key_type 									metadata_type;
		typedef Node_CItr 	       							node_const_iterator;
		typedef Node_Itr 									node_iterator;
		typedef typename node_const_iterator::value_type 	const_iterator;
		typedef typename node_iterator::value_type 			iterator;
		
		key_type
		interval_sum(key_const_reference, key_const_reference);

	private:
		typedef typename base_type::const_reference 	const_reference;
		typedef typename base_type::const_pointer 		const_pointer;

		typedef typename _Alloc::template rebind<metadata_type>::other __rebind_m;
		typedef typename __rebind_m::const_reference       metadata_const_reference;
		typedef typename __rebind_m::reference 		metadata_reference;

		virtual node_const_iterator
		node_begin() const = 0;

		virtual node_iterator
		node_begin() = 0;

		virtual node_const_iterator
		node_end() const = 0;

		virtual node_iterator
		node_end() = 0;

		virtual cmp_fn&
		get_cmp_fn() = 0;
		
		key_type
		prefix_sum(key_const_reference);
		
	protected:
		inline void
		operator()(node_iterator, node_const_iterator) const;
};

#define BELONG interval_sum_node_update<Node_CItr, Node_Itr, Cmp_Fn, _Alloc>

TTT
void BELONG::
operator ()(node_iterator cur, node_const_iterator end) const
{
	node_iterator l = cur.get_l_child(), r = cur.get_r_child();
	int sl = l == end ? 0 : l.get_metadata(),
		sr = r == end ? 0 : r.get_metadata();
	const_cast<metadata_reference>(cur.get_metadata()) = **cur + sl + sr;
}

TTT
typename BELONG::key_type BELONG::
prefix_sum(key_const_reference key)
{
	key_type ret = 0;
	cmp_fn C = get_cmp_fn();
	node_iterator i = node_begin();
	for(node_iterator l, r; i != node_end(); )
	{
		l = i.get_l_child(), r = i.get_r_child();
		if(C(key, **i)) i = l;
		else
		{
			ret += **i;
			if(l != node_end()) ret += l.get_metadata();
			i = r;
		}
	}
	return ret;
}

TTT
typename BELONG::key_type BELONG::
interval_sum(key_const_reference l, key_const_reference r)
{
	return prefix_sum(r) - prefix_sum(l - 1);
}

int main()
{
	tree<int, null_type, less<int>, splay_tree_tag, interval_sum_node_update> tr;
	int n;
	scanf("%d", &n);
	for(int i = 1; i <= n; i++)
	{
		int t;
		scanf("%d", &t);
		tr.insert(t);
	}
	printf("%d\n", tr.interval_sum(1, n));
}