手打 splay_tree 模板类 支持 find, rank, 前驱后继, iterator

Posted 2020-12-05 js2xxx

• 迭代器中序遍历整棵树，++或--为前驱后继
• 可重复插入，查找，排名（<或<=），删除，大小，相同元素个数
• 功能超全的 stl-like splay 平衡树，摈弃平板电视，还不快收入囊中！
• 插入删除的方式为new/delete，不喜勿喷~

如有错请指正！（支持版本：C++11及以上）

#include <algorithm>

using namespace std;

template<typename T>
class splay_tree_iterator;

template<typename T, typename _Less = less<T>, typename _Equ = equal_to<T>>
class splay_tree
{
    public:
        friend class splay_tree_iterator<T>;
        typedef _Less lesscmp;
        typedef _Equ equcmp;
        typedef splay_tree_iterator<T> iterator;

        splay_tree() : _r(nullptr) { }

        iterator root() { return iterator(_r); }

        iterator begin()
        {
            _node *p = _r;
            if(!p) return iterator(nullptr);
            while(p->s[0]) p = p->s[0];
            return iterator(p);
        }

        iterator end() { iterator k = _end(); return ++k; }

        void insert(T o)
        {
            _node *k = _insert(nullptr, _r, o);
            if(k != _r) _splay(k);
        }

        iterator find(T o)
        {
            _node *k = _find(_r, o);
            if(!k) return end();
            return iterator(k);
        }

        void erase(T o)
        {
            _node *k = _find(_r, o);
            if(!k) return;
            _splay(k);
            if(k->cnt > 1) { --k->cnt, --k->size; return; }
            _node *p = k->_prec();
            if(p)
            {
                _splay(p), p->s[1] = k->s[1];
                if(p->s[1]) p->s[1]->f = p;
                delete k;
            }
            else if((p = k->_succ())) _splay(p), delete k;
            else delete k, _r = nullptr;
            k = nullptr;
            if(p) _sync(p);
        }
        
        size_t size() { return _r ? _r->size : 0; }
        
        size_t rank(T o, bool equ = false) { return _rank(_r, o, equ); }

    private:
        lesscmp _cl;
        equcmp _ce;
        struct _node
        {
            T v;
            int cnt, size;
            _node *f, *s[2];
            _node * _succ()
            {
                _node *k = this;
                if(!k->s[1])
                {
                    for(; k->f && k->f->s[1] == k; k = k->f);
                    return k->f;
                }
                else
                {
                    for(k = k->s[1]; k->s[0]; k = k->s[0]);
                    return k;
                }
            }
            _node * _prec()
            {
                _node *k = this;
                if(!k->s[0])
                {
                    for(; k->f && k->f->s[0] == k; k = k->f);
                    return k->f;
                }
                else
                {
                    for(k = k->s[0]; k->s[1]; k = k->s[1]);
                    return k;
                }
            }
        } *_r;

        inline void _sync(_node *k)
        {
            k->size = (k->s[0] ? k->s[0]->size : 0) + (k->s[1] ? k->s[1]->size : 0) + k->cnt;
        }

        inline void _rotate(_node *k)
        {
            _node *l = k->f, *m = l->f;
            bool isr = l->s[1] == k;

            if(m) m->s[m->s[1] == l] = k;
            k->f = m;

            l->s[isr] = k->s[!isr], k->s[!isr] = l, l->f = k;
            if(l->s[isr]) l->s[isr]->f = l;
            _sync(l), _sync(k);
        }

        void _splay(_node *k)
        {
            while(k->f)
            {
                _node *l {k->f}, *m {l->f};
                if(!m) _rotate(k);
                else if(m->s[1] == l ^ l->s[1] == k)
                    _rotate(k), _rotate(k);
                else _rotate(l), _rotate(k);
            }
            _r = k;
        }

        _node * _insert(_node *f, _node *&k, T o)
        {
            if(!k) return k = new _node { o, 1, 1, f, nullptr, nullptr };
            ++k->size;
            if(_cl(k->v, o)) return _insert(k, k->s[1], o);
            if(!_ce(k->v, o)) return _insert(k, k->s[0], o);
            ++k->cnt;
            return k;
        }

        _node *_find(_node *k, T o)
        {
            if(!k) return nullptr;
            if(_cl(k->v, o)) return _find(k->s[1], o);
            if(!_ce(k->v, o)) return _find(k->s[0], o);
            return k;
        }

        iterator _end()
        {
            _node *p = _r;
            if(!p) return nullptr;
            while(p->s[1]) p = p->s[1];
            return iterator(p);
        }
        
        size_t _rank(_node *k, T o, bool equ)
        {
            if(!k) return 0;
            if(_cl(k->v, o)) return _rank(k->s[1], o, equ) + (k->s[0] ? k->s[0]->size : 0) + k->cnt;
            if(!_ce(k->v, o)) return _rank(k->s[0], o, equ);
            return (k->s[0] ? k->s[0]->size : 0) + (equ ? k->cnt : 1);
        }
};

template<typename T>
class splay_tree_iterator : iterator<bidirectional_iterator_tag, T>
{
    public:
        friend class splay_tree<T>;
        splay_tree_iterator(const splay_tree_iterator<T> &i) : _n(i._n), _p(i._p) { }
        splay_tree_iterator &operator++()
        {
            _p = _n, _n = _n->_succ();
            return *this;
        }
        splay_tree_iterator operator++(int)
        {
            splay_tree_iterator p(*this);
            _p = _n, _n = _n->_succ();
            return p;
        }
        splay_tree_iterator &operator--()
        {
            if(!_n) _n = _p, _p = nullptr;
            else _n = _n->_prec();
            return *this;
        }
        splay_tree_iterator operator--(int)
        {
            splay_tree_iterator p(*this);
            if(!_n) _n = _p, _p = nullptr;
            else _n = _n->_prec();
            return p;
        }
        T &operator *() { return _n->v; }
        T *operator ->() { return &(_n->v); }
        size_t count() { return _n->cnt; }
        bool operator != (const splay_tree_iterator<T> &i) { return this->_n != i._n; }
        bool operator == (const splay_tree_iterator<T> &i) { return !(*this != i); }
    private:
        typedef typename splay_tree<T>::_node ptr;
        splay_tree_iterator(ptr *p) : _n(p), _p(nullptr) { }
        ptr *_n, *_p;
};