splay_tree.h

#ifndef SPLAY_TREE_H
#define SPLAY_TREE_H
 
#ifdef ENABLE_TREE_VISUALIZATION
#include "../../visualization/tree_visual/tree_visualization.h"
#endif
 
#ifdef __cplusplus
#include <functional>
#include <iostream>
#include <memory>
#include <queue>
#include <vector>
#endif

template <typename T> class splay_tree {
  private:
    struct node {
        T info;
        std::shared_ptr<node> right;
        std::shared_ptr<node> left;
        node(T key = 0) : info(key), right(nullptr), left(nullptr) {}
    };
    std::shared_ptr<node> root;
    size_t _size{0};
 
  public:
    inline explicit splay_tree(std::vector<T> v = {}) noexcept : root(nullptr) {
        if (!v.empty()) {
            for (T& x : v) {
                this->insert(x);
            }
        }
    }

    inline explicit splay_tree(const splay_tree& s) : root(s.root), _size(s._size) {}

    inline splay_tree& operator=(const splay_tree& s) {
        root = s.root;
        _size = s._size;
        return *this;
    }
 
    inline ~splay_tree() { root = nullptr; }

    inline void clear() {
        root = nullptr;
        _size = 0;
    }

    inline void insert(T key) {
        root = _insert(root, key);
        _size++;
    }

    inline void remove(T key) {
        root = _remove(root, key);
        _size--;
    }

    inline bool search(T key) {
        std::shared_ptr<node> ans = splay(root, key);
        return (ans && ans->info == key);
    }

    inline size_t size() { return _size; }
 
    class Iterator;
 
    inline Iterator begin() {
        std::vector<T> ino = this->inorder();
        return Iterator(0, ino);
    }
 
    inline Iterator end() {
        std::vector<T> ino = this->inorder();
        return Iterator(ino.size(), ino);
    }

    inline std::vector<T> inorder() {
        std::vector<T> path;
        _inorder([&](std::shared_ptr<node> callbacked) { path.push_back(callbacked->info); }, root);
        return path;
    }

    inline std::vector<T> preorder() {
        std::vector<T> path;
        _preorder([&](std::shared_ptr<node> callbacked) { path.push_back(callbacked->info); },
                  root);
        return path;
    }

    inline std::vector<T> postorder() {
        std::vector<T> path;
        _postorder([&](std::shared_ptr<node> callbacked) { path.push_back(callbacked->info); },
                   root);
        return path;
    }

    inline std::vector<std::vector<T>> level_order() {
        std::vector<std::vector<T>> path;
        std::queue<std::shared_ptr<node>> q;
        q.push(root);
        while (!q.empty()) {
            size_t size = q.size();
            std::vector<T> level;
            for (size_t i = 0; i < size; i++) {
                std::shared_ptr<node> current = q.front();
                q.pop();
                level.push_back(current->info);
                if (current->left) {
                    q.push(current->left);
                }
                if (current->right) {
                    q.push(current->right);
                }
            }
            path.push_back(level);
        }
        return path;
    }

#ifdef TREE_VISUALIZATION_H
    inline void visualize() {
        std::string _generated = generate_visualization();
        tree_visualization::visualize(_generated);
    }
#endif

    inline friend std::ostream& operator<<(std::ostream& out, splay_tree<T>& t) {
        std::vector<T> order = t.inorder();
        for (int i = 0; i < order.size(); i++) {
            if (i != order.size() - 1) {
                out << order[i] << ", ";
            } else {
                out << order[i] << '\n';
            }
        }
        return out;
    }
 
  private:
    std::shared_ptr<node> rrotate(std::shared_ptr<node> _node) {
        std::shared_ptr<node> y = _node->left;
        _node->left = y->right;
        y->right = _node;
        return y;
    }
 
    std::shared_ptr<node> lrotate(std::shared_ptr<node> _node) {
        std::shared_ptr<node> y = _node->right;
        _node->right = y->left;
        y->left = _node;
        return y;
    }
 
    std::shared_ptr<node> splay(std::shared_ptr<node> _node, T key) {
        if (!_node || _node->info == key) {
            return _node;
        }
        if (_node->info > key) {
            if (_node->left == NULL)
                return _node;
            if (_node->left->info > key) {
                _node->left->left = splay(_node->left->left, key);
                _node = rrotate(_node);
            } else if (_node->left->info < key) {
                _node->left->right = splay(_node->left->right, key);
                if (_node->left->right != NULL)
                    _node->left = lrotate(_node->left);
            }
            return (_node->left == NULL) ? _node : rrotate(_node);
        } else {
            if (_node->right == NULL)
                return _node;
            if (_node->right->info > key) {
                _node->right->left = splay(_node->right->left, key);
                if (_node->right->left != NULL)
                    _node->right = rrotate(_node->right);
            } else if (_node->right->info < key) {
                _node->right->right = splay(_node->right->right, key);
                _node = lrotate(_node);
            }
            return (_node->right == NULL) ? _node : lrotate(_node);
        }
    }
 
    std::shared_ptr<node> _insert(std::shared_ptr<node> root, T key) {
        if (!root) {
            std::shared_ptr<node> nn = std::make_shared<node>(key);
            return nn;
        }
        root = splay(root, key);
        if (root->info == key) {
            return root;
        }
 
        std::shared_ptr<node> nn = std::make_shared<node>(key);
        if (root->info > key) {
            nn->right = root;
            nn->left = root->left;
            root->left = nullptr;
        } else {
            nn->left = root;
            nn->right = root->right;
            root->right = nullptr;
        }
        return nn;
    }
 
    std::shared_ptr<node> _remove(std::shared_ptr<node> root, T key) {
        if (!root) {
            return nullptr;
        }
        root = splay(root, key);
        if (key != root->info) {
            return root;
        }
 
        std::shared_ptr<node> temp;
        temp = root;
        if (!root->left) {
            root = root->right;
        } else {
            root = splay(root->left, key);
            root->right = temp->right;
        }
        return root;
    }
 
    void _inorder(std::function<void(std::shared_ptr<node>)> callback, std::shared_ptr<node> root) {
        if (root) {
            _inorder(callback, root->left);
            callback(root);
            _inorder(callback, root->right);
        }
    }
 
    void _postorder(std::function<void(std::shared_ptr<node>)> callback,
                    std::shared_ptr<node> root) {
        if (root) {
            _postorder(callback, root->left);
            _postorder(callback, root->right);
            callback(root);
        }
    }
 
    void _preorder(std::function<void(std::shared_ptr<node>)> callback,
                   std::shared_ptr<node> root) {
        if (root) {
            callback(root);
            _preorder(callback, root->left);
            _preorder(callback, root->right);
        }
    }
 
    std::string generate_visualization() {
        std::string _generate = _inorder_gen(root);
        return _generate;
    }
 
    std::string _inorder_gen(std::shared_ptr<node> root) {
        std::string _s;
        if (std::is_same_v<T, char> || std::is_same_v<T, std::string>) {
            if (root->left) {
                _s += root->info;
                _s += "->";
                _s += root->left->info;
                _s += "\n";
                _s += _inorder_gen(root->left);
            }
            if (root->right) {
                _s += root->info;
                _s += "->";
                _s += root->right->info;
                _s += "\n";
                _s += _inorder_gen(root->right);
            }
        } else {
            if (root->left) {
                _s += std::to_string(root->info) + "->" + std::to_string(root->left->info) + "\n" +
                      _inorder_gen(root->left);
            }
            if (root->right) {
                _s += std::to_string(root->info) + "->" + std::to_string(root->right->info) + "\n" +
                      _inorder_gen(root->right);
            }
        }
        return _s;
    }
};

template <typename T> class splay_tree<T>::Iterator {
  private:
    std::vector<T> elements;
    int64_t index;
 
  public:
    explicit Iterator(const int64_t& index, std::vector<T>& els) noexcept
        : index(index), elements(els) {}

    Iterator& operator=(int64_t index) {
        this->index = index;
        return *(this);
    }

    Iterator& operator++() {
        if (this->index < elements.size()) {
            this->index++;
        }
        return *(this);
    }

    Iterator operator++(int) {
        Iterator it = *this;
        ++*(this);
        return it;
    }

    Iterator& operator--() {
        if (this->index > 0) {
            this->index--;
        }
        return *(this);
    }

    Iterator operator--(int) {
        Iterator it = *this;
        --*(this);
        return it;
    }

    bool operator!=(const Iterator& it) { return index != it.index; }

    T operator*() { return elements[index]; }
};
 
#endif