avl/avl.cpp

#include <fstream>

#include <cstdint>


// #define BALANCE_IN_POINTERS
// #define NODE_REUSE

#define _max(a,b) (((a) > (b)) ? (a) : (b)) //mhhhm speedy
// #define _max(a,b) ((a) > (b)) * (a) + ((b) > (a)) * (b)
typedef unsigned long uintptr_t
#define get_parent(node) ((Node *)((uintptr_t)(node)->parent & ~0x3))

#ifdef BALANCE_IN_POINTERS
  #define get_bf(node) (((int)((uintptr_t)(node)->parent & 0x3)) - 1 )
  #define get_balance(node) ((node)?(get_bf(node)):(0))
#else
  #define get_bf(node) ((node)->balance)
  #define get_balance(node) ((node) ? get_bf(node) : (0))
#endif

inline int _abs(int n) {
    int mask = n >> ((sizeof(int)*8) -1);
    return (mask + n)^mask;
}

template <typename T>
class PointerStack {
public:
    T* array;
    size_t capacity;
    size_t top;

    PointerStack(size_t size) {
        capacity = size;
        array = size ? new T[capacity] : nullptr;
        top = -1;
    }

    ~PointerStack() {
        delete[] array;
    }

    inline void push(T value) {
        if (isFull()) {
            size_t newCapacity = 2 * capacity;
            T *newArray = new T[newCapacity];

            for (size_t i = 0; i <= top; ++i) {
                newArray[i] = array[i];
            }

            capacity = newCapacity;
            delete[] array;

            array = newArray;
        }

        array[++top] = value;
    }

    inline T pop() {
        if (empty()) {
            throw std::out_of_range("pop from empty stack");
        }

        return array[top--];
    }

    inline const T& peek() const {
        if (empty()) {
            return -1;
        }

        return array[top];
    }

    inline bool empty() const {
        return top == size_t(-1);
    }

    inline bool isFull() const {
        return top == capacity - 1;
    }

    inline void flush () {
      if (empty())
        return;
      for (size_t i = 0; i <= top; ++i) {
        delete array[i];
      }
    }
};

template <typename T>
struct Tree {
public:
  Tree(
    #ifdef NODE_REUSE
      size_t stacksize = 0
    #endif
    ) : root(nullptr), treeSize(0)
    #ifdef NODE_REUSE
      , stack(stacksize) 
    #endif 
    {}
  ~Tree() { 
    delete root;
    #ifdef NODE_REUSE
      if (stack.capacity)
        stack.flush();
    #endif
  }

  struct Node {
    Node(const T& data, Node* parent = nullptr) : data(data), left(nullptr), right(nullptr), parent(nullptr)
    #ifndef BALANCE_IN_POINTERS
    , balance(0)
    #endif
    {}

    ~Node() {
      delete left;
      delete right;
    }

    T data;
    Node* left;
    Node* right;
    Node* parent;
    #ifndef BALANCE_IN_POINTERS
    int balance;
    #endif
    #ifndef __PROGTEST__
    void debugDotShow(std::ofstream &ofs) {
      if (left) {
        ofs << data << " -> " << left->data << '\n';
        left->debugDotShow(ofs);
      }
      if (right) {
        ofs << data << " -> " << right->data << '\n';
        right->debugDotShow(ofs);
      }
    }
    #endif
  };

  Node* root;
  size_t treeSize;
  #ifdef NODE_REUSE
    PointerStack<Node *> stack;
  #endif

  size_t size() const { return treeSize; }
  const T* find(const T& value) const { 
    Node* i = root;
    while (i) {
      if (i->data == value)
        break;
      i = i->data < value ? i->right : i->left;
    }
    return i ? &i->data : nullptr;
  }

  bool erase(const T& value) { 
    Node* i = root;
    while (i) {
      if (i->data == value)
        break;
      i = i->data < value ? i->right : i->left;
    }

    if (!i)
      return false;
    
    remove(i);

    return true;
  }

  #ifndef __PROGTEST__
  void debugDotShow() {
    if (!root)
      return;
    std::ofstream ofs;
    ofs.open("viz.dot");
    ofs << "digraph Tree {\n";
    root->debugDotShow(ofs);
    ofs << "}";

    ofs.close();
    system("cat viz.dot | dot -Tx11 &");

  }

    void testParent(Node *n) const {
      if (!n)
        return;
      if (n->left) {
        if (n->left->parent != n)
          throw std::logic_error(std::to_string(n->left->data) + " parent error");
        if (n->left == n)
          throw std::logic_error(std::to_string(n->data) + " selfchild");
        testParent(n->left);
      }
      if (n->right) {
        if (n->right->parent != n)
          throw std::logic_error(std::to_string(n->right->data) + " parent error");
        if (n->right == n)
          throw std::logic_error(std::to_string(n->data) + " selfchild");
        testParent(n->right);
      }
    }
  #endif

  static inline void set_bf(Node* n, int bf) {
    #ifdef BALANCE_IN_POINTERS
    n->parent = (Node *)((uintptr_t)get_parent(n) | (uintptr_t)(bf+1));
    #else
    n->balance = bf;
    #endif
  }
  static inline void set_parent(Node* n, Node* parent) {
    #ifdef BALANCE_IN_POINTERS
    n->parent = (Node*)((uintptr_t)parent | ((uintptr_t)n->parent & 0x3));
    #else
    n->parent = parent;
    #endif
  }

  bool insert(T value) {
    Node* inserted;
    Node* parent = nullptr;
    Node* current = root;
    int bf, bfOld;

    while (current) {
      parent = current;

      if (value == current->data)
        return false;

      current = value < current->data ? current->left : current->right;
    }

    inserted = getEmptyNode(value);
    set_parent(inserted, parent);
    set_bf(inserted, 0);
    ++treeSize;
    
    if (parent) {
      if (inserted->data < parent->data)
        parent->left = inserted;
      else
        parent->right = inserted;
    } else {
      root = inserted;
    }

    bf = 0;
    while(inserted) {
            parent = get_parent(inserted);

            if (parent) {
                bfOld = get_bf(inserted);

                if (parent->right == inserted) {
                    inserted = balanceTree(inserted, bf);
                    parent->right = inserted;
                }else {
                    inserted = balanceTree(inserted, bf);
                    parent->left = inserted;
                }

                // parent bf
                if (!inserted->left && !inserted->right) {
                    // leaf node
                    if (parent->left == inserted) bf = -1;
                    else bf = 1;
                } else {
                    // index node
                    bf = 0;
                    if (_abs(bfOld) < _abs(get_bf(inserted))) {
                        if (parent->left == inserted) bf = -1;
                        else bf = 1;
                    }
                }

            } else if(inserted == root){
                root = balanceTree(root, bf);
                break;
            }
            if (bf == 0) break;

            inserted = parent;
        }
    return true;
  }

  static inline Node* balanceTree(Node* n, int bf) {
    int bfChild;
    int heightDiff = get_balance(n) + bf;

    if (n) {
        if(heightDiff < -1 && n->left) {
            // balance left
            if(get_balance(n->left) <= 0) {
                bfChild = get_bf(n->left);
                n = LL(n, heightDiff, &bfChild, nullptr);
                set_bf(n, bfChild);
            } else {
                n = LR(n, heightDiff);
            }
        } else if(heightDiff > 1 && n->right) {
            // balance right
            if(get_balance(n->right) >= 0) {
                bfChild = get_bf(n->right);
                n = RR(n, heightDiff, &bfChild, nullptr);
                set_bf(n, bfChild);
            } else {
                n = RL(n, heightDiff);
            }
        } else {
            set_bf(n, get_bf(n) + bf);
        }
    }
    return n;
  }
  
  static inline Node* LL(Node* parent, int bfParent, int* bfChild, int* heightDiff) {
    int parentRight, leftChild, rightChild;
    Node *child = parent->left;

    leftChild = (child->left)?(1):(0);
    rightChild = (child->right)?(1):(0);
    if (*bfChild < 0) {
        // child->left > child->right
        leftChild = rightChild - (*bfChild);
        parentRight = leftChild + 1 + bfParent;
        if (heightDiff)
            *heightDiff = _max(leftChild, _max(rightChild, parentRight)+1) - (leftChild + 1);

    } else {
        // child->left <= child->right
        rightChild = leftChild + (*bfChild);
        parentRight = rightChild + 1 + bfParent;
        if (heightDiff)
            *heightDiff = _max(leftChild, _max(rightChild, parentRight)+1) - (rightChild + 1);
    }
    *bfChild = (_max(rightChild, parentRight) + 1) - leftChild;
    set_bf(parent, parentRight - rightChild);

    parent->left = child->right;
    if (child->right)
        set_parent(child->right, parent);
    child->right = parent;
    set_parent(child, get_parent(parent));
    set_parent(parent, child);

    return child;
  }

  static inline Node* RR(Node* parent, int bfParent, int* bfChild, int* heightDiff) {
    int parentLeft, childLeft, childRight;
    Node *child = parent->right;

    childLeft = (child->left)?(1):(0);
    childRight = (child->right)?(1):(0);
    if (*bfChild < 0) {
        // child->left > child->right
        childLeft = childRight - (*bfChild);
        parentLeft = childLeft + 1 - bfParent;
        if (heightDiff)
            *heightDiff = _max(childRight, _max(childLeft, parentLeft)+1) - (childLeft + 1);

    } else {
        // child->left <= child->right
        childRight = childLeft + (*bfChild);
        parentLeft = childRight + 1 - bfParent;
        if (heightDiff)
            *heightDiff = _max(childRight, _max(childLeft, parentLeft)+1) - (childRight + 1);

    }
    *bfChild = childRight - (_max(childLeft, parentLeft) + 1);
    set_bf(parent, childLeft - parentLeft);

    parent->right = child->left;
    if (child->left)
        set_parent(child->left, parent);
    child->left = parent;
    set_parent(child, get_parent(parent));
    set_parent(parent, child);

    return child;
  }

  static inline Node* RL(Node* parent, int bfParent) {
    int bfChild, height_delta = 0;
    Node* child = parent->right;
    Node* ret;

    if (child->left) {
        bfChild = get_bf(child->left);
        parent->right = LL(child, get_bf(child), &bfChild, &height_delta);
    } else {
        bfChild = get_bf(child);
    }

    ret = RR(parent, bfParent+height_delta, &bfChild, nullptr);
    set_bf(ret, bfChild);
    return ret;
  }

  static inline Node* LR(Node* parent, int bfParent) {
    int bfChild, height_delta = 0;
    Node* child = parent->left;
    Node* ret;

    if (child->right) {
        bfChild = get_bf(child->right);
        parent->left = RR(child, get_bf(child), &bfChild, &height_delta);
    } else {
        bfChild = get_bf(child);
    }

    ret = LL(parent, bfParent-height_delta, &bfChild, nullptr);
    set_bf(ret, bfChild);
    return ret;
  }

  inline void remove(Node *node) {
    if (!node)
      return;

    Tree rightSubtree;
    Node* p = nullptr;
    Node* current;
    Node* swap = nullptr;
    int bf = 0, bfOld;

    rightSubtree.root = node->right;
    swap = rightSubtree.findSmolest(rightSubtree.root);

    if (swap) {
        // next exists
        if (get_parent(swap)) {
            if (get_parent(swap) != node) {
                get_parent(swap)->left = swap->right;
                if (swap->right)
                    set_parent(swap->right, get_parent(swap));
            }
        }
        if (get_parent(node)) {
            // replace node by next
            if (get_parent(node)->left == node) {
                get_parent(node)->left = swap;
            } else {
                get_parent(node)->right = swap;
            }
        }

        // re-link pointers
        if (node->right != swap) {
            swap->right = node->right;
            if (node->right) set_parent(node->right, swap);
            current = get_parent(swap);
            bf = 1;
        }else{
            current = swap;
            bf = -1;
        }

        swap->left = node->left;
        if (node->left) set_parent(node->left, swap);
        set_parent(swap, get_parent(node));

        // inherit nodes balance factor
        set_bf(swap, get_bf(node));

    } else {
        // there's no right subTree
        p = get_parent(node);
        if (p) {
            if (p->left == node) {
                p->left = node->left;
                bf = 1;
            } else {
                p->right = node->left;
                bf = -1;
            }
        }
        if (node->left)
            set_parent(node->left, p);

        current = get_parent(node);
    }

    // reset root
    if (root == node) {
        root = swap;
        if (!swap) {
            if (node->left) root = node->left;
        }
    }

    // balancing
    while(current) {
        p = get_parent(current);
        if (p) {
            // if parent exists
            bfOld = get_bf(current);

            if (p->right == current) {
                current = balanceTree(current, bf);
                p->right = current;
            }else {
                current = balanceTree(current, bf);
                p->left = current;
            }

            // bf for parent
            if (!current->left && !current->right) {
                // leaf
                if (p->left == current) bf = 1;
                else bf = -1;
            } else {
                // index
                bf = 0;
                if (_abs(bfOld) > _abs(get_bf(current))) {
                    if (p->left == current) bf = 1;
                    else bf = -1;
                }
            }

        } else if(current == root){
            root = balanceTree(root, bf);
            break;
        }
        if (bf == 0) break;

        current = p;
    }
    rightSubtree.root = nullptr;
    node->left = nullptr;
    node->right = nullptr;
    // delete node;
    saveEmptyNode(node);
    --treeSize;
  }

  inline Node* findSmolest(Node *n) {
    if (!n)
      return nullptr;
    Node *i = n;
    for (;i->left; i = i->left);
    return i;
  }

  inline void saveEmptyNode(Node *n) {
    #ifdef NODE_REUSE
    stack.push(n);
    #else
    delete n;
    #endif
  }

  inline Node* getEmptyNode(T value) {
    #ifdef NODE_REUSE
    if (stack.empty())
      return new Node(value);
    auto ret = stack.pop();
    ret->data = value;
    return ret;
    #else
    return new Node(value);
    #endif
  }

};