rbtree.cc

#include <iostream>
using namespace std;

enum Color { BLACK, RED };


template <typename T>
class Node_ {
public:
    Node_ () { };
    Node_ (Color c, Node_<T>* l, T t, Node_<T>* r) : color(c), left(l), root(t), right(r) { };
    Node_ (const Node_<T> &other) : color(other.color), left(other.left), root(other.root), right(other.right) { };
    ~Node_ () { };

    Color color;
    Node_<T>* left;
    T root;
    Node_<T>* right;

    void reassign(Color, Node_<T>*, T, Node_<T>*);
};

template <typename T>
void Node_<T>::reassign(Color c, Node_<T>* l, T t, Node_<T>* r) {
    color = c;
    left = l;
    root = t;
    right = r;
}


template <typename T>
struct Triad_ {
    Node_<T>* parent;
    Node_<T>* lsub;
    Node_<T>* rsub;
};

template<typename T>
class Tree_ {
public:
    Tree_ () { root = NULL; };
    Tree_ (const Tree_<T> &other) { root = other.root; };
    ~Tree_ () { _teardown(root); };
    void insert(T x) { root = _insert(x, root); };
    void print_tree() { _print_tree(root); };
private:
    Node_<T>* root;
    Node_<T>* balanced(Triad_<T>& nodes, Node_<T>* ll, T lt, Node_<T>* lr, T t, Node_<T>* rl, T rt, Node_<T>* rr);
    Node_<T>* balance(Node_<T>* tree);
    Node_<T>* ins(T x, Node_<T>* t);    
    Node_<T>* _insert(T x, Node_<T>* tree);
    void _print_tree(Node_<T>* tree);
    void _teardown(Node_<T>* tree);
};

template<typename T>
Node_<T>* Tree_<T>::balanced(Triad_<T>& nodes, Node_<T>* ll, T lt, Node_<T>* lr, T t,
			     Node_<T>* rl, T rt, Node_<T>* rr) {
    nodes.lsub->reassign(BLACK, ll, lt, lr);
    nodes.rsub->reassign(BLACK, rl, rt, rr);
    nodes.parent->reassign(RED,nodes.lsub, t, nodes.rsub);
    return nodes.parent;
}

template<typename T>
Node_<T>* Tree_<T>::balance(Node_<T>* tree) {
    // The lines in balance are intentionally long to help with visualization.
    if (tree->color == BLACK && tree->left && tree->left->color == RED) {
	if (tree->left->left && tree->left->left->color == RED) {
	    Triad_<T> nodes = {tree, tree->left->left, tree->left};
	    return balanced(nodes,
			    tree->left->left->left, tree->left->left->root, tree->left->left->right,
			    tree->left->root,
			    tree->left->right, tree->root, tree->right);
	}

	if (tree->left->right && tree->left->right->color == RED) {
	    Triad_<T> nodes = {tree, tree->left, tree->left->right};
	    return balanced(nodes,
			    tree->left->left, tree->left->root,
			    tree->left->right->left, tree->left->right->root, tree->left->right->right,
			    tree->root, tree->right);
	}
    }

    if (tree->color == BLACK && tree->right && tree->right->color == RED) {
	if (tree->right->left && tree->right->left->color == RED) {
	    Triad_<T> nodes = {tree, tree->right, tree->right->left};
	    return balanced(nodes,
			    tree->left, tree->root,
			    tree->right->left->left, tree->right->left->root, tree->right->left->right,
			    tree->right->root, tree->right->right);
	}

        if (tree->right->right && tree->right->right->color == RED) {
	    Triad_<T> nodes = {tree, tree->right, tree->right->right};
            return balanced(nodes,
			    tree->left, tree->root,
			    tree->right->left, tree->right->root,
			    tree->right->right->left, tree->right->right->root, tree->right->right->right);
	}
    }
    return tree;
}


template<typename T>
Node_<T>* Tree_<T>::ins(T x, Node_<T>* t) {
    if (!t)
	return new Node_<T>(RED, NULL, x, NULL);
	
    if (x < t->root) {
	t->left = ins(x, t->left);
	return balance(t);
    }
    else if (t->root < x) {
	t->right = ins(x, t->right);
	return balance(t);
    }
    else
	return t;
}

template<typename T>
Node_<T>* Tree_<T>::_insert(T x, Node_<T>* tree) {
    Node_<T>* bal_tree = ins(x, tree);
    bal_tree->color = BLACK;
    return bal_tree;
}

template<typename T>
void Tree_<T>::_print_tree(Node_<T>* tree) {
    if (tree->color == RED)
	cout << tree->root << "[style=filled, fillcolor=red];" << endl;
    else
	cout << tree->root << "[style=filled, fillcolor=black];" << endl;

    if (tree->left) {
	cout << tree->root << " -> " << tree->left->root << ";" << endl;
	_print_tree(tree->left);
    }

    if (tree->right) {
	cout << tree->root << " -> " << tree->right->root << ";" << endl;
	_print_tree(tree->right);
    }	
}

template<typename T>
void Tree_<T>::_teardown(Node_<T>* tree){
    if (tree->left)
	_teardown(tree->left);
    if (tree->right)
	_teardown(tree->right);
    delete tree;
}


typedef Tree_<int> Tree;


int main() {
    Tree  tree;
    int input[] = {15, 5, 6, 17, 2, 19, 16, 8, 14, 1, 12, 7, 3, 18, 0, 11, 4, 9, 13, 10};
    int i = 0, size = sizeof(input) / sizeof(int);

    for (i=0; i < size; i++) {
	tree.insert(input[i]);
    }

    cout << "digraph tree { " << endl;
    cout << "node [fontcolor=white, nodesep=1];" << endl;
    tree.print_tree();
    cout << "}" << endl;

    return 0;
}