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 {
public:
  typedef _Key Key;

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

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);
  }

  void replaceChild(Self* child, Self* replace)
  {
    assert(child == left || child == right);

    if(child == left)
      left = replace;
    else if(child == right)
      right = replace;

    if(replace)
      replace->parent = this;

    child->parent = nullptr;

    recalc();
  }

public:
  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);

    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;
  //   else if(node->parent->right == node)
  //     node->parent->right = replace;

  //   if(replace)
  //     {
  //       replace->parent = node->parent;
  //       replace->recalc();
  //     }
  // }

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

    Node* child;

    if(leaf->left == nullptr && leaf->right == nullptr)
      child = nullptr;
    else if(leaf->left == nullptr)
      child = leaf->right;
    else if(leaf->right == nullptr)
      child = leaf->left;
    else
      assert(false);

    if(leaf->parent)
      {
        leaf->parent->replaceChild(leaf, child);
        rebalance(leaf->parent, 1);
      }
    else
      {
        root = child;
        if(child)
          child->parent = nullptr;
      }

    // if(leaf->left == nullptr && leaf->right == nullptr)
    //   {
    //     if(leaf->parent != nullptr)
    //       {
    //         replaceInParent(leaf, nullptr);
    //         rebalance(leaf->parent, 1, 2);
    //       }
    //     else
    //       root = nullptr;
    //   }
    // else if(leaf->left == nullptr)
    //   {
    //     if(leaf->parent != nullptr)
    //       {
    //         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)
  {
    for(Node* iter = start; iter != nullptr; iter = iter->parent)
      {
        iter->recalc();

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

        if(absBalance == 2) // needs rotation
          {
            // 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;
          }
        else if(absBalance == abortCond) // can 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;
        partner->parent = nullptr;
      }
    else
      node->parent->replaceChild(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;
        partner->parent = nullptr;
      }
    else
      node->parent->replaceChild(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);

    Node* parent = nullptr;

    if(root == nullptr)
      // Tree is empty, add node as root
      root = node;
    else
      {
        // else, find a suitable parent
        parent = findByKey(node->key);
        assert(parent);

        if(node->key < parent->key)
          parent->left = node;
        else if(node->key > parent->key)
          parent->right = node;
        else // Node with same key exists
          return false;
      }

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

    // rebalance tree
    if(parent)
      rebalance(parent,0);

    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->parent == nullptr)
      {
        root = replacement;
        replacement->parent = nullptr;
      }
    else
      node->parent->replaceChild(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->recalc();
  }

  Node* lookup(Key key) const
  {
    if(root == nullptr)
      return nullptr;

    Node* node = findByKey(key);

    if(node == nullptr)
      __builtin_unreachable();

    if(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
};