one-file-projects/bintree.h

#include <cassert>

#ifdef DUMP
#include <iostream>
using std::cout;
using std::endl;
#endif

typedef unsigned int Umword;
typedef int Mword;

template <typename Key>
class Bin_tree;

template <typename _Key = Umword>
struct Bin_tree_node {
private:
  friend class Bin_tree<_Key>;

  void recalc()
  {
    Umword leftHeight =
      (left == nullptr) ? 0 : left->height + 1;
    Umword rightHeight =
      (right == nullptr) ? 0 : right->height + 1;

    balance = leftHeight - rightHeight;
    height = (leftHeight < rightHeight ? rightHeight : leftHeight);
  }

public:
  typedef _Key Key;

protected:
  typedef Bin_tree_node<Key> Self;
  typedef Bin_tree<Key> Tree;

  constexpr Bin_tree_node(Key key)
    : key(key), parent(nullptr), left(nullptr),
      right(nullptr), tree(nullptr),
      balance(0), height(0)
  {}

  constexpr Bin_tree_node()
    : Bin_tree_node(0)
  {}

  Self* parent;
  Self* left;
  Self* right;
  Tree* tree;

  // AVL Tree
  Mword balance;
  Umword height;

public:
  Key key;
};

template <typename Key = Umword>
class Bin_tree {
protected:
  typedef Bin_tree_node<Key> Node;
private:
  Node* root;

  Node* findByKey(Key key) const
  {
    Node* iterparent = nullptr;
    Node* iter = root;

    while(iter != nullptr)
      {
        iterparent = iter;

        if(key < iter->key) {
          iter = iter->left;
        } else if(key > iter->key) {
          iter = iter->right;
        } else {
          break;
        }
      }

    return iterparent;
  }

  Node* findLeft(Node* start) const
  {
    assert(start != nullptr);
    Node* iter = start;
    while(iter->left != nullptr)
      {
        iter = iter->left;
      }
    return iter;
  }

  void replaceInParent(Node* node, Node* replace) {
    assert(node->parent);
    assert(node->parent->left == node
           || node->parent->right == node);

    assert(node->parent);

    if(node->parent->left == node)
      {
        node->parent->left = replace;
        if(replace) replace->parent = node->parent;
      }
    else if(node->parent->right == node)
      {
        node->parent->right = replace;
        if(replace) replace->parent = node->parent;
      }
  }

  void leafRemove(Node* leaf)
  {
    assert(leaf->left == nullptr
           || leaf->right == nullptr);

    if(leaf->left == nullptr && leaf->right == nullptr)
      {
        if(leaf->parent != nullptr)
          {
            assert(leaf->parent->left == leaf
                   || leaf->parent->right == leaf);
            replaceInParent(leaf, nullptr);
            rebalance(leaf->parent, 1, 2);
          }
        else
          root = nullptr;
      }
    else if(leaf->left == nullptr)
      {
        if(leaf->parent != nullptr)
          {
            assert(leaf->parent->left == leaf
                   || leaf->parent->right == leaf);
            replaceInParent(leaf, leaf->right);
            rebalance(leaf->parent, 1, 2);
          }
        else
          {
            root = leaf->right;
            leaf->right->parent = nullptr;
          }
      }
    else if(leaf->right == nullptr)
      {
        if(leaf->parent != nullptr)
          {
            replaceInParent(leaf, leaf->left);
            rebalance(leaf->parent, 1, 2);
          }
        else
          {
            root = leaf->left;
            leaf->left->parent = nullptr;
          }
      }

    leaf->parent = nullptr;
    leaf->left = nullptr;
    leaf->right = nullptr;
    leaf->tree = nullptr;

    leaf->recalc();
  }

  void
  rebalance(Node* start, Mword abortCond, Mword rebalanceCond)
  {
    for(Node* iter = start; iter != nullptr; iter = iter->parent)
      {
        iter->recalc();

        Mword absBalance =
          (iter->balance < 0) ? -(iter->balance) : iter->balance;

        if(absBalance == abortCond)
          {
            break;
          }
        else if(absBalance == rebalanceCond)
          {
            // Left side is taller
            if(iter->balance > 0)
              {
                // Left side's right side is taller
                if(iter->left && iter->left->balance < 0)
                  {
                    rotateLeft(iter->left);
                  }
                rotateRight(iter);
              }
            // Right side is taller
            else
              {
                // Right side's left side is taller
                if(iter->right && iter->right->balance > 0)
                  {
                    rotateRight(iter->right);
                  }
                rotateLeft(iter);
              }

            // then abort
            break;
          }
      }
  }

  void rotateLeft(Node* node) {
    assert(node);
    assert(node->right);

    Node* partner = node->right;

    node->right = partner->left;

    if(node->right)
      node->right->parent = node;

    // No parent, partner is the new root
    if(node->parent == nullptr)
      root = partner;
    else
      replaceInParent(node, partner);

    partner->parent = node->parent;
    node->parent = partner;

    partner->left = node;

    assert(node->parent == nullptr
           || node->parent->left == node
           || node->parent->right == node);

    assert(partner->parent == nullptr
           || partner->parent->left == partner
           || partner->parent->right == partner);

    assert(node->left == nullptr
           || node->left->parent == node);

    assert(node->right == nullptr
           || node->right->parent == node);

    assert(partner->left == nullptr
           || partner->left->parent == partner);

    assert(partner->right == nullptr
           || partner->right->parent == partner);

    node->recalc();
    partner->recalc();
  }

  void rotateRight(Node* node) {
    assert(node);
    assert(node->left);

    Node* partner = node->left;

    node->left = partner->right;

    if(node->left)
      node->left->parent = node;

    // No parent, partner is the new root
    if(node->parent == nullptr)
      root = partner;
    else
      replaceInParent(node,partner);

    partner->parent = node->parent;
    node->parent = partner;

    partner->right = node;

    assert(node->parent == nullptr
           || node->parent->left == node
           || node->parent->right == node);

    assert(partner->parent == nullptr
           || partner->parent->left == partner
           || partner->parent->right == partner);

    assert(node->left == nullptr
           || node->left->parent == node);

    assert(node->right == nullptr
           || node->right->parent == node);

    assert(partner->left == nullptr
           || partner->left->parent == partner);

    assert(partner->right == nullptr
           || partner->right->parent == partner);

    node->recalc();
    partner->recalc();

  }

  #ifdef DUMP
  void dump(Node* node, int indent = 0)
  {
    if(node == nullptr) return;
    assert(node->tree == this);

    assert(node->left == nullptr
           || node->left->parent == node);
    dump(node->left, indent+1);
    for(int i = 0; i < indent; i++)
      cout << "   ";
    cout << node->key << endl;

    assert(node->right == nullptr
           || node->right->parent == node);
    dump(node->right, indent+1);
  }
  #endif

public:
  constexpr Bin_tree() : root(nullptr)
  {}

  bool insert(Node* node)
  {
    assert(node->tree == nullptr);

    if(root == nullptr)
      {
        root = node;
        root->parent = nullptr;
        root->left = nullptr;
        root->right = nullptr;
        root->tree = this;
        return true;
      }

    Node* parent = findByKey(node->key);

    if(node->key < parent->key)
      {
        parent->left = node;
      }
    else if(node->key > parent->key)
      {
        parent->right = node;
      }
    else
      {
        return false;
      }

    node->parent = parent;
    node->left = nullptr;
    node->right = nullptr;
    node->tree = this;
    node->height = 0;
    node->balance = 0;

    rebalance(parent,0,2);

    return true;
  }

  void remove(Node* node)
  {
    assert(node);
    assert(node->tree == this);

    if(node->left == nullptr || node->right == nullptr)
      {
        leafRemove(node);
        return;
      }

    Node* replacement = findLeft(node->right);

    leafRemove(replacement);

    if(node == root)
      {
        root = replacement;
        replacement->parent = nullptr;
      }
    else
      replaceInParent(node, replacement);

    replacement->left = node->left;
    replacement->right = node->right;
    replacement->tree = this;

    if(node->left)
      node->left->parent = replacement;

    if(node->right)
      node->right->parent = replacement;

    replacement->recalc();

    node->left = nullptr;
    node->right = nullptr;
    node->parent = nullptr;
    node->tree = nullptr;
  }

  Node* lookup(Key key) const {
    Node* node = findByKey(key);

    if(node == nullptr || node->key != key)
      return nullptr;

    return node;
  }

  #ifdef DUMP
  void dump()
  {
    cout << "---------------------------------------------" << endl;
    dump(root);
    cout << "---------------------------------------------" << endl;
  }
  #endif
};

template<typename T, typename Key = typename T::Key>
class Bin_tree_t : Bin_tree<Key> {
private:
  typedef Bin_tree<Key> Base;
public:
  inline
  bool insert(T* node)
  { return Base::insert(static_cast<typename Base::Node*>(node)); }

  inline
  void remove(T* node)
  { Base::remove(static_cast<typename Base::Node*>(node)); }

  inline
  T* lookup(Key key)
  { return static_cast<T*>(Base::lookup(key)); }

  #ifdef DUMP
  inline
  void dump()
  { Base::dump(); }
  #endif
};