BBST/splay_tree/splay_tree_map_ushi.hpp

#include"./splay_tree_base.hpp"

template<typename T,typename E,typename F>
struct splay_tree_map_ushi_node{
    using np=splay_tree_map_ushi_node*;
    np ch[2]={0,0};
    int sz=1;
    T key;E val;E sum;F op;
    splay_tree_map_ushi_node(T key,E val,F op):key(key),val(val),sum(val),op(op){}
    int size(np t){return t?t->sz:0;}
    np push(){}
    np update(){
        sz=size(ch[0])+1+size(ch[1]);
        sum=val;
        if(ch[0])sum=op(ch[0]->sum,sum);
        if(ch[1])sum=op(sum,ch[1]->sum);
        return this;
    }
};

template<typename T,typename E,typename F>
struct splay_tree_map_ushi:splay_tree_base<splay_tree_map_ushi_node<T,E,F>>{
    using node=splay_tree_map_ushi_node<T,E,F>;
    using super=splay_tree_base<node>;
    using np=node*;
    F op;
    splay_tree_map_ushi(F op=F()):op(op){}
    inline int size(){return super::size(super::root);}
    int get_idx(T key){return super::lower_bound([&](np t){return t->key<key;});}
    inline void insert(T key,E val){
        return super::insert(get_idx(key),new node(key,val,op));
    }
    inline E back(){return (*this)[size()-1];}
    inline E front(){return (*this)[0];}
    bool contains(T key){
        int idx=get_idx(key);
        return idx<size()&&super::get(idx)->key==key;
    }
    E operator[](T key){return super::get(get_idx(key))->val;}
    E fold_idx(int l,int r){
        auto [p,q]=super::split(super::root,l);
        auto [s,t]=super::split(q,r-l);
        T res=s?s->sum:E();
        super::root=super::merge(p,super::merge(s,t));
        return res;
    }
    E fold(T l,T r){
        int ll=get_idx(l);
        int rr=get_idx(r);
        auto [p,q]=super::split(super::root,ll);
        auto [s,t]=super::split(q,rr-ll);
        E res=s?s->sum:E();
        super::root=super::merge(p,super::merge(s,t));
        return res;
    }
    template<typename G>
    void update(T key,G g){
        int idx=super::lower_bound([&](np t){return t->key<key;});
        np t=super::get(idx);
        g(t->val);
        t->update();
    }
};

#line 1 "BBST/splay_tree/splay_tree_base.hpp"
#include<vector>
#include<tuple>
#include<assert.h>

template<typename Node>
struct splay_tree_base{
    using np=Node*;
    np root=0;
    inline int size(np t){return t?t->sz:0;}
    int size(){return size(root);}
    np rot(np t,bool b){
        np s=t->ch[1-b];
        assert(s);
        np u=s->ch[b];
        t->ch[1-b]=u;
        s->ch[b]=t;
        t->update();s->update();
        return s;
    }
    void erase(int idx){
        auto [p,q]=split(root,idx);
        auto [s,t]=split(q,1);
        delete s;
        root=merge(p,t);
    }
    void insert(int idx,np val){
        if(size(root)==idx){
            root=merge(root,val);
        }else{
            auto [p,q]=split(root,idx);
            root=merge(merge(p,val),q);
        }
    }
    np merge(np s,np t){
        if(!s||!t)return s?s:t;
        s=splay(s,size(s)-1);
        s->ch[1]=t;
        return s->update();
    }
    std::pair<np,np> split(np t,int idx){
        if(size(t)==idx)return std::make_pair(t,nullptr);
        t=splay(t,idx);
        np s=t->ch[0];
        t->ch[0]=nullptr;
        return std::make_pair(s,t->update());
    }
    np splay(np t,int idx){
        if(idx==size(t->ch[0]))return t;

        bool b1=size(t->ch[0])<idx;
        if(b1)idx-=size(t->ch[0])+1;

        if(idx==size(t->ch[b1]->ch[0]))return rot(t,1-b1);
        bool b2=size(t->ch[b1]->ch[0])<idx;
        if(b2)idx-=size(t->ch[b1]->ch[0])+1;

        t->ch[b1]->ch[b2]=splay(t->ch[b1]->ch[b2],idx);
        if(b1==b2)t=rot(t,1-b2);
        else t->ch[b1]=rot(t->ch[b1],1-b2);

        return rot(t,1-b1);
    }
    std::vector<np> output(){
        std::vector<np>res;
        auto dfs=[&](auto dfs,np t)->void{
            if(!t)return;
            dfs(dfs,t->ch[0]);
            res.push_back(t);
            dfs(dfs,t->ch[1]);
        };
        dfs(dfs,root);
        return res;
    }
    np get_root(){return root;}
    np get(int i){root=splay(root,i);return root;}
    
    template<typename C>
    int lower_bound(C less){
        int res=__lower_bound(root,less);
        if(res<size())root=splay(root,res);
        return res;
    }
    template<typename C>
    int __lower_bound(np t,C less){
        if(!t)return 0;
        bool b=less(t);
        if(b)return size(t->ch[0])+1+__lower_bound(t->ch[1],less);
        else return __lower_bound(t->ch[0],less);
    }
};
#line 2 "BBST/splay_tree/splay_tree_map_ushi.hpp"

template<typename T,typename E,typename F>
struct splay_tree_map_ushi_node{
    using np=splay_tree_map_ushi_node*;
    np ch[2]={0,0};
    int sz=1;
    T key;E val;E sum;F op;
    splay_tree_map_ushi_node(T key,E val,F op):key(key),val(val),sum(val),op(op){}
    int size(np t){return t?t->sz:0;}
    np push(){}
    np update(){
        sz=size(ch[0])+1+size(ch[1]);
        sum=val;
        if(ch[0])sum=op(ch[0]->sum,sum);
        if(ch[1])sum=op(sum,ch[1]->sum);
        return this;
    }
};

template<typename T,typename E,typename F>
struct splay_tree_map_ushi:splay_tree_base<splay_tree_map_ushi_node<T,E,F>>{
    using node=splay_tree_map_ushi_node<T,E,F>;
    using super=splay_tree_base<node>;
    using np=node*;
    F op;
    splay_tree_map_ushi(F op=F()):op(op){}
    inline int size(){return super::size(super::root);}
    int get_idx(T key){return super::lower_bound([&](np t){return t->key<key;});}
    inline void insert(T key,E val){
        return super::insert(get_idx(key),new node(key,val,op));
    }
    inline E back(){return (*this)[size()-1];}
    inline E front(){return (*this)[0];}
    bool contains(T key){
        int idx=get_idx(key);
        return idx<size()&&super::get(idx)->key==key;
    }
    E operator[](T key){return super::get(get_idx(key))->val;}
    E fold_idx(int l,int r){
        auto [p,q]=super::split(super::root,l);
        auto [s,t]=super::split(q,r-l);
        T res=s?s->sum:E();
        super::root=super::merge(p,super::merge(s,t));
        return res;
    }
    E fold(T l,T r){
        int ll=get_idx(l);
        int rr=get_idx(r);
        auto [p,q]=super::split(super::root,ll);
        auto [s,t]=super::split(q,rr-ll);
        E res=s?s->sum:E();
        super::root=super::merge(p,super::merge(s,t));
        return res;
    }
    template<typename G>
    void update(T key,G g){
        int idx=super::lower_bound([&](np t){return t->key<key;});
        np t=super::get(idx);
        g(t->val);
        t->update();
    }
};