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

转载于:https://www.cnblogs.com/js2xxx/p/9446963.html