avl_tree.h

#ifndef AVL_TREE_H
#define AVL_TREE_H
 
#ifdef __cplusplus
#include <functional>
#include <memory>
#include <queue>
#include <string>
#include <vector>
#endif
 
template <typename T> class avl_tree {
public:
  explicit avl_tree(std::vector<T> _elements = {}) noexcept : root(nullptr) {
    if (!_elements.empty()) {
      for (T &x : _elements) {
        this->insert(x);
      }
    }
  }
 
  explicit avl_tree(const avl_tree &a) noexcept : root(a.root), _size(a._size) {
 
 
  }
 
  avl_tree &operator=(const avl_tree &a) noexcept {
    root = a.root;
    _size = a._size;
    return *this;
  }
 
  ~avl_tree() noexcept {}
 
  void insert(T key) {
    root = _insert(root, key);
    _size++;
  }
 
  void clear() {
    root = nullptr;
    _size = 0;
    return;
  }
 
  T get_root() const { return this->root->info; }
 
  bool search(T key) { return _search(root, key); }
 
  class Iterator;
 
  Iterator begin() {
    std::vector<T> ino = this->inorder();
    return Iterator(0, ino);
  }
 
  Iterator end() {
    std::vector<T> ino = this->inorder();
    return Iterator(ino.size(), ino);
  }
 
  size_t size() const { return _size; }
  void remove(T key) {
    root = _remove(root, key);
    _size--;
  }
 
  std::vector<T> inorder() const {
    std::vector<T> path;
    _inorder(
      [&](std::shared_ptr<node> callbacked) {
        path.push_back(callbacked->info);
      },
      root);
    return path;
  }
  std::vector<T> preorder() const {
    std::vector<T> path;
    _preorder(
      [&](std::shared_ptr<node> callbacked) {
        path.push_back(callbacked->info);
      },
      root);
    return path;
  }
  std::vector<T> postorder() const {
    std::vector<T> path;
    _postorder(
      [&](std::shared_ptr<node> callbacked) {
        path.push_back(callbacked->info);
      },
      root);
    return path;
  }
 
  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;
  }
 
  friend std::ostream & operator << (std::ostream &out, avl_tree<T> &t){
    std::vector<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:
  typedef struct node {
    T info;
    int64_t height{0};
    std::shared_ptr<node> left;
    std::shared_ptr<node> right;
    node(T key) : info(key), left(nullptr), right(nullptr) {}
  } node;
 
  std::shared_ptr<node> root;
  size_t _size{};
 
  int64_t height(std::shared_ptr<node> root) {
    if (root == nullptr)
      return 0;
    return 1 + std::max(height(root->left), height(root->right));
  }
 
  std::shared_ptr<node> createNode(T info) {
    std::shared_ptr<node> nn = std::make_shared<node>(info);
    return nn;
  }
 
  int64_t getBalance(std::shared_ptr<node> root) {
    return height(root->left) - height(root->right);
  }
 
  std::shared_ptr<node> rightRotate(std::shared_ptr<node> root) {
    std::shared_ptr<node> t = root->left;
    std::shared_ptr<node> u = t->right;
    t->right = root;
    root->left = u;
    return t;
  }
 
  std::shared_ptr<node> leftRotate(std::shared_ptr<node> root) {
    std::shared_ptr<node> t = root->right;
    std::shared_ptr<node> u = t->left;
    t->left = root;
    root->right = u;
    return t;
  }
 
  std::shared_ptr<node> minValue(std::shared_ptr<node> root) {
    if (root->left == nullptr)
      return root;
    return minValue(root->left);
  }
 
  std::shared_ptr<node> _insert(std::shared_ptr<node> root, T item) {
    std::shared_ptr<node> nn = createNode(item);
    if (root == nullptr)
      return nn;
    if (item < root->info) {
      root->left = _insert(root->left, item);
    }
    else if (item > root->info) {
      root->right = _insert(root->right, item);
    }
    else {
      return root;
    }
    int b = getBalance(root);
    if (b > 1) {
      if (getBalance(root->left) < 0)
        root->left = leftRotate(root->left);
      return rightRotate(root);
    } else if (b < -1) {
      if (getBalance(root->right) > 0)
        root->right = rightRotate(root->right);
      return leftRotate(root);
    }
    return root;
  }
 
  std::shared_ptr<node> _remove(std::shared_ptr<node> root, T key) {
    if (root == nullptr)
      return root;
    if (key < root->info)
      root->left = _remove(root->left, key);
    else if (key > root->info)
      root->right = _remove(root->right, key);
 
    else {
      if (!root->right) {
        std::shared_ptr<node> temp = root->left;
        root = nullptr;
        return temp;
      } else if (!root->left) {
        std::shared_ptr<node> temp = root->right;
        root = nullptr;
        return temp;
      }
      std::shared_ptr<node> temp = minValue(root->right);
      root->info = temp->info;
      root->right = _remove(root->right, temp->info);
    }
    return root;
  }
 
  bool _search(std::shared_ptr<node> root, T key) {
    while (root) {
      if (root->info < key) {
        root = root->right;
      } else if (root->info > key) {
        root = root->left;
      } else {
        return true;
      }
    }
    return false;
  }
 
  void _inorder(std::function<void(std::shared_ptr<node>)> callback,
                std::shared_ptr<node> root) const {
    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) const {
    if (root) {
      _inorder(callback, root->left);
      _inorder(callback, root->right);
      callback(root);
    }
  }
 
  void _preorder(std::function<void(std::shared_ptr<node>)> callback,
                 std::shared_ptr<node> root) const {
    if (root) {
      callback(root);
      _inorder(callback, root->left);
      _inorder(callback, root->right);
    }
  }
 
};
 
template <typename T> class avl_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 && elements[index] != it.elements[it.index];
  }
 
  T operator*() { return elements[index]; }
};
 
#endif