An AVL tree implementation

.gitignore

tree
*.o
*.swp

main.cpp

#include <cstdlib>
#include <ctime>
#include <unistd.h>

#include <iostream>
using std::cout; using std::endl;

#include "tree.hpp"

int main()
{

   //const int testcase[] = {66, 20, 43, 87, 97, 40, 20, 19, 78, 50, 97, 30, 92, 67, 77, 1, 87};
   const int testlength = 128;//sizeof(testcase)/sizeof(int); // 32

   srand(getpid() * time(0));

   tree test;
   int n;

   for(int i = 0; i < testlength; ++i){
      //n = testcase[i];//rand() % 100;
      n = rand() % 100;
      cout << n << ' ';
      test.insert(n, 0);
   }
   cout << endl;

   test.inorderPrint();
   
   return 0;
}

Makefile

cpp_flags := -Wall -Wextra -Weffc++ -Wshadow -pedantic
cpp_flags := $(cpp_flags) -ggdb
#cpp_flags := $(cpp_flags) -O3
#cpp_flags := $(cpp_flags) -std=c++0x

obj_files := tree.o node.o main.o

exe_file := tree

.PHONY: all clean

all: $(exe_file)

$(exe_file): $(obj_files)
	g++ $(cpp_flags) $^ -o $@

clean:
	@rm -f $(obj_files) $(exe_file)

%.o: %.cpp
	g++ $(cpp_flags) -c $^ -o $@

node.cpp

#include <iostream>
using std::cout;
using std::endl;

#include <algorithm>
using std::max;

#include <cassert>
// assert

#include "tree.hpp"
#include "node.hpp"

tree::node::node(tree::key_t k, tree::value_t v)
   :
      parent(NULL), right(NULL), left(NULL),
      key(k), value(v), _height(1)
{

}

tree::node * tree::node::getRoot(){
   node * r = this;

   while(r->parent) r = r->parent;

   return r;
}

size_t tree::node::lheight(){
   return left  ? left->height()  : 0;
}

size_t tree::node::rheight(){
   return right ? right->height() : 0;
}

size_t tree::node::height(){
   return _height;
}

size_t tree::node::realHeight(){
   if(left && right){
      return 1 + max(left->realHeight(), right->realHeight());
   }else if(left || right){
      return 1 + (left ? left : right)->realHeight();
   }else{
      return 1;
   }
}

bool tree::node::isBalanced(){
   size_t lh = lheight();
   size_t rh = rheight();
   return (rh+2>lh) && (lh+2>rh);
}

bool tree::node::isBST(){
   key_t dummy;
   bool b = true;
   return isBST(dummy, b);
}

bool tree::node::isBST(key_t & prevK, bool & isFirst){
   if (left && !left->isBST(prevK, isFirst)) return false;

   if(isFirst){
      isFirst = false;
   }else{
      if (key < prevK) return false;
   }
   prevK = key;

   if (right && !right->isBST(prevK, isFirst)) return false;

   return true;

}

// recursive tree insert
void tree::node::insert(node * n){

   node *& target = n->key < key ? left : right;

   if(target)
      target->insert(n);
   else{
      target = n;
      n->parent = this;
      updateHeight();
   }

   // the entire tree should be a balanced BST after each
   // insertion
   assert(getRoot()->isBST());
   assert(getRoot()->isBalanced());
}

void tree::node::updateHeight(){
   size_t startHeight = _height;

   _height = 1 + max(lheight(), rheight());

   if (!isBalanced()) balance();

   if (parent && startHeight != _height) parent->updateHeight();
}

void tree::node::balance(){
   // do nothing if we're already balanced
   if (!isBalanced()){

      // TODO - give proofs for height changes in seperate
      // text file. It'd be too long and annoying to have them here.

      size_t lh, rh;

      if(lheight() > rheight()){
         /* RIGHT ROTATION */

         lh = left->lheight();
         rh = left->rheight();

         if(lh > rh){         // Left-Heavy
            _height -= 2;

         }else if(lh == rh){  // Perfectly balanced
            _height -= 1;
            left->_height += 1;

         }else /*(lh < rh)*/{ // Right-Heavy
            _height -= 2;
            left->_height -= 1;
            left->right->_height += 1;
            left->rotateLeft();

         }
         rotateRight();

      }else{
         /* LEFT ROTATION */

         lh = right->lheight();
         rh = right->rheight();

         if(rh > lh){         // Right-Heavy
            _height -= 2;

         }else if(rh == lh){  // Perfectly balanced
            _height -= 1;
            right->_height += 1;

         }else /*(rh < lh)*/{ // Left-Heavy
            _height -= 2;
            right->_height -= 1;
            right->left->_height += 1;
            right->rotateRight();

         }
         rotateLeft();
      }

      // at this point, this node should
      // be balanced. ALWAYS.
      assert(isBalanced());

      // Also, at this point, the three nodes involved in the rotation
      // should all have the correct heights.
      assert(parent->height() == parent->realHeight());
      assert(parent->left->height() == parent->left->realHeight());
      assert(parent->right->height() == parent->right->realHeight());

   }else{
      assert(false && "unecessary call to tree::node::balance");
   }
}

// ROTATIONS
//
// schematics for the left and right rotations are shown below
// `o` represents an arbitrary balanced subtree
// `C` represents a particular subtree of `N`
// `P` represents the parent of `T` (either on the right or
// on the left, it doesn't matter)
// `T` represents the current node
// `N` represents the node to be rotate into `T`'s place
//
// In both cases, there are 3 bi-directional links that must be
// modified, for a total of 6 pointers that must be changed.
//
// I think the clearest way to accomplish this might be to
// make a generic function that does the re-linking and make
// rotateRight and rotateLeft wrappers that call it.
//
/**************************************

   -------- RIGHT -------------
    
       P              P
       |              |
       T              N
      / \     =>     / \
     N   o          o   T
    / \                / \
   o   C              C   o

   ---------- LEFT -------------
                                           
      P                   P
      |                   |
      T                   N
     / \      =>         / \
    o   N               T   o
       / \             / \
      C   o           o   C
******************************************/

// P, N, C and T are as depicted in the diagram above
//
// _out pointers are the appropriate outgoing members
//  (either `right` or `left`)
//
//  in pointers are the nodes themselves
//
//  rotate(P, T, T_child, N_child);
void tree::node::rotate(
      node * const P,
      node * const T,
      node * &     T_child,
      node * &     N_child
      ){

   // acts as a writable NULL, similar to /dev/null
   node * dummy;

   // get references to relevant variables
   node *& P_child = P ? 
      (P->left == T ? P->left : P->right)
      : dummy;

   node * N = T_child;
   node * C = N_child;

   // re-assign links

   P_child = N;

   N->parent = P;
   N_child = T;

   T->parent = N;
   T_child = C;

   if (C) C->parent = T;
}

void tree::node::rotateRight(){
   rotate(parent, this, this->left, this->left->right);
}

void tree::node::rotateLeft(){
   rotate(parent, this, this->right, this->right->left);
}

// TODO - DEBUGGING FUNCT, REMOVE
void tree::node::inorderPrint(int i){
   if(left) left->inorderPrint(i+1);
   cout << key << ' ';
   if(right) right->inorderPrint(i+1);
   if(!left && !right) cout << "LN(" << i << ") ";
}

node.hpp

#ifndef NODE_HPP
#define NODE_HPP

#include "tree.hpp"

class tree::node {

   private:
      node * parent;
      node * right;
      node * left;

      tree::key_t   key;
      tree::value_t value;

      size_t _height;
      size_t lheight();
      size_t rheight();
      size_t height();
      size_t realHeight();

      bool isBalanced();
      bool isBST();
      bool isBST(key_t &, bool &);

      void updateHeight();
      void balance();

      void rotate(node*const, node*const, node*&, node*&);

      void rotateRight();
      void rotateLeft();

   public:
      node(key_t, value_t);

      void insert(node *);

      void inorderPrint(int i);

      node * getRoot();
};

#endif

test.bash

#!/bin/bash

if [ ! -f tree ]
then
   echo "Executable (tree) does not exist, aborting"
   exit 1
fi

if [ -f treetest ]
then
   echo "tempfile (treetest) exists already, aborting"
   exit 2
fi

./tree > treetest;
count=1

# set up trap
ctrl_c() {
   echo "completed $count runs"
   rm treetest
   exit 0
}

trap ctrl_c SIGINT
# end setup trap

while [ $? -eq 0 ]
do
   ./tree > treetest
   count=$((++i))
done

cat treetest

echo "completed $count runs"

rm treetest

todo.txt

Shit that needs to get done!

test.bash
=========

- parameterize file names
- generate random filename for temp file based on pid and timestamp or something

node.cpp
========

- add deletions
- move from OO to procedural style, use tree class as wrapper for private procedural manipulations
   - This just seems like a more natural way of expressing things. Rotations and balancing are seriously
     non-intuitive, as the relative position of `this` changes in the process.

tree.cpp
========

- templatize
- Move to wrapping procedural rather than OO code (see above)
- Add iterators (pre-order, post-order, inorder, r_inorder) (output and input)
   - Is it possible to make these bi-directional? That would be interesting...
- add lookup/retrieval operations

Makefile / Build process
========================

- Add object dir
- Add optional `debug` and `release` targets
   - wrap assetions and debugging functions in `#ifdef DEBUG`'s and pass -DDEBUG in `debug` target

Misc.
=====

- Make proofs.txt, containing descriptions of the more complicated operations and how they work
- Check if it's worth actually open-sourcing this and making it into something real
   - I've seen people looking for stl-style tree containers before.. Do they exist somewhere else?

tree.cpp

#include <iostream>
using std::cout;
using std::endl;

#include "tree.hpp"
#include "node.hpp"

tree::tree() : root(NULL) {}

void tree::insert(key_t k, value_t v){

   node * n = new node(k,v);

   if(root){
      root->insert(n);

      // In the balancing process, `root` may cease to be
      // the real root of the tree. In this case we have
      // to re-acquire the true root of the tree.
      root = root->getRoot();
   }else{
      root = n;
   }
}

void tree::inorderPrint(){
   root->inorderPrint(1);
   cout << endl;
}

tree.hpp

#ifndef TREE_H
#define TREE_H


class tree {
   // TYPES
   private:
      struct node;
   public:
      typedef int value_t;
      typedef int key_t;

   // MEMBERS
   private:

      node * root;

   public:

      tree();

      void insert(key_t, value_t);
      void inorderPrint();
};


#endif