AVL Tree

avl_tree.hpp

//
//  avl_tree2.hpp
//  AVLTree
//
//  Created by Seiji Emery on 11/18/14.
//  Copyright (c) 2014 Seiji Emery. All rights reserved.
//

#ifndef AVLTree_avl_tree2_hpp
#define AVLTree_avl_tree2_hpp

#include <algorithm>
#include <functional>
#include <cassert>
#include <array>

template <typename T, bool hasSetProperty = false>
class AVLTree {
protected:
    struct Node {
        std::array<Node*, 2> children;
        T value;
        int8_t depth = 0;
        
        Node (const T & value, Node * left = nullptr, Node * right = nullptr)
            : children({ left, right }), value(value) { recalcDepth(); }
        ~Node () {
            if (children[0] != nullptr)
                delete children[0];
            if (children[1] != nullptr)
                delete children[1];
        }
        
        void recalcDepth () {
            depth = std::max(children[0] ? children[0]->depth : 0,
                             children[1] ? children[1]->depth : 0) + 1;
        }
        
        static bool isValidChild (int child) {
            return child == 0 || child == 1;
        }
        
        static int other (int child) {
            assert(isValidChild(child));
            return child ? 0 : 1;
        }
        
        int select (const T & val) {
            return val < value ? 0 : 1;
        }
        
        static Node * deepcopy (const Node * node) {
            return node ? new Node (node->value, deepcopy(node->children[0]), deepcopy(node->children[1]))
                        : nullptr;
        }
        
        Node * rotateLR (int left, int right) {
            assert(isValidChild(left) && isValidChild(right) && left != right);
            assert(children[left] != nullptr);
            
            Node * tmp = children[left]->children[right];
            this->children[left]->children[right] = tmp->children[left];
            tmp->children[left] = children[left];
            this->children[left] = tmp->children[right];
            tmp->children[right] = this;
            tmp->children[left]->recalcDepth();
            tmp->children[right]->recalcDepth();
            tmp->recalcDepth();
            return tmp;
        }
        
        Node * rotateLL (int left, int right) {
            Node * tmp = children[left];
            children[left] = tmp->children[right];
            tmp->children[right] = this;
            tmp->children[right]->recalcDepth();
            tmp->recalcDepth();
            return tmp;
        }

        Node * insert (const T & val) {
            if (hasSetProperty) {   // Don't allow inserting the same value multiple times.
                if (val == value)   // Makes implementing sets much more efficient, but not much other use...
                    return this;    // (hence defaulted-off template parameter)
            }
            int left = select(val), right = other(left);
            if (children[left] == nullptr) {
                children[left] = new Node(val);
                if (depth < 2)
                    depth = 2;
            } else {
                children[left] = children[left]->insert(val);
                recalcDepth();
                
                if (children[left]->depth > (children[right] ? children[right]-> depth : 0) + 1) {
                    // rebalance
                    
                    if (children[left]->children[left] == nullptr) {
                        assert(children[left]->children[right] != nullptr);
                        return rotateLR(left, right);
                    } else {
                        return rotateLL(left, right);
                    }
                }
            }
            return this;
        }
        
        int balanceFactor () const {
            return (children[0] ? children[0]->depth : 0) - (children[1] ? children[1]->depth : 0);
        }

        
        Node * rebalance () {
            int factor = balanceFactor();
            
            if (abs(factor) >= 2) {
                int dir = (factor > 0) ? 0 : 1;
                
                if (children[dir]->children[dir] == nullptr) {
                    assert(children[dir]->children[other(dir)] != nullptr);
                    return rotateLR(dir, other(dir));
                } else {
                    return rotateLL(dir, other(dir));
                }
            }
            return this;
        }
        
        Node * recursivelyRebalanceSubtree (int dir) {
            assert(isValidChild(dir));
            
            if (children[dir] != nullptr) {
                children[dir] = children[dir]->recursivelyRebalanceSubtree(dir);
                return this->rebalance();
            }
            return this;
        }
        
        Node * remove (const T & val) {
            if (val == value) {
                int dir = (balanceFactor() >= 0 ? 0 : 1), rdir = other(dir);
                Node * repl = children[dir];
                Node * parent = nullptr;
                
                if (repl == nullptr) {
                    // node has no children
                    assert(children[rdir] == nullptr);
                    delete this;
                    return nullptr;
                }
                
                while (repl->children[rdir] != nullptr) {
                    parent = repl;
                    repl = repl->children[rdir];
                }
                
                std::swap(value, repl->value);
                if (parent != nullptr) {
                    // repl is somewhere deep in the subtree
                    parent->children[rdir] = repl->children[dir];
                    repl->children[dir] = nullptr;
                    assert(repl->children[rdir] == nullptr);
                    parent->recalcDepth();
                    children[dir] = children[dir]->recursivelyRebalanceSubtree(rdir);
                } else {
                    // repl is node's direct child
                    children[dir] = repl->children[dir];
                    recalcDepth();
                }
                delete repl;
                return this;
            } else {
                int dir = select(val);
                if (children[dir] != nullptr) {
                    children[dir] = children[dir]->remove(val);
                    return rebalance();
                }
                return this;
            }
        }
        
        void traverse (const std::function<void(T &)> & callback) {
            if (children[0] != nullptr)
                children[0]->traverse(callback);
            callback(value);
            if (children[1] != nullptr)
                children[1]->traverse(callback);
        }
        
        void traverse (const std::function<void(const T &)> & callback) const {
            if (children[0] != nullptr)
                children[0]->traverse(callback);
            callback(value);
            if (children[1] != nullptr)
                children[1]->traverse(callback);
        }
        
        T * find (const T & val) {
            if (val == this->value) {
                return &value;
            } else {
                int left = select(val);
                if (children[left] != nullptr)
                    return children[left]->find(val);
                return nullptr;
            }
        }
        
        static void print (const Node * node, std::ostream & out, int level) {
            assert(level > 0);
            for (auto i = 0; i < level; ++i)
                out << '\t';
            if (node != nullptr) {
                out << "Node: " << node->value << " (depth " << static_cast<int>(node->depth) << ")\n";
                print(node->children[0], out, level+1);
                print(node->children[1], out, level+1);
            } else {
                out << "empty (depth 0)\n";
            }
        }
    };
    Node * root = nullptr;
public:
    AVLTree () {}
    AVLTree (const std::initializer_list<T> & values) {
        for (const auto & val : values) {
            insert(val);
        }
    }
    AVLTree (const AVLTree & tree)
        : root(Node::deepcopy(tree.root))
    {}
    
    AVLTree (AVLTree && other)
    : root(other.root)
    {
        other.root = nullptr;
    }
    ~AVLTree () {
        if (root != nullptr)
            delete root;
    }
    
    T * find (const T & value) {
        if (root)
            return root->find(value);
        return nullptr;
    }
    
    const T * find (const T & value) const {
        if (root)
            return root->find(value);
        return nullptr;
    }
    
    void remove (const T & value) {
        if (root) {
            root = root->remove(value);
        }
    }
    
    void insert (const T & value) {
        if (root == nullptr) {
            root = new Node(value);
        } else {
            root = root->insert(value);
        }
    }
    
    void print (std::ostream & out) const {
        out << "AVL Tree\n";
        Node::print(root, out, 1);
    }
    
    void each (const std::function<void (const T &)> & callback) const {
        if (root != nullptr)
            root->traverse(callback);
    }
    void each (const std::function<void (T &)> & callback) {
        if (root != nullptr)
            root->traverse(callback);
    }
    
    friend std::ostream & operator << (std::ostream & os, const AVLTree & tree) {
        return tree.print(os), os;
    }
};

#endif

main.cpp

//
//  main.cpp
//  AVLTree
//
//  Created by Seiji Emery on 11/18/14.
//  Copyright (c) 2014 Seiji Emery. All rights reserved.
//

#include <iostream>
#include "avl_tree2.hpp"

template <typename T>
void printTree(const std::initializer_list<T> & args) {
    (AVLTree<T> { args }).print(std::cout);
}

template <typename T>
void printHasMember (const AVLTree<T> & tree, const T & value) {
    std::cout << (tree.find(value) != nullptr ? "tree contains " : "tree does not contain ") << value << '\n';
}

template <typename T>
void printRemoval (AVLTree<T> & tree, const T & value) {
    std::cout << "Removing " << value << '\n';
    tree.remove(value);
    tree.print(std::cout);
}

template <typename T>
void printInsertion (AVLTree<T> & tree, const T & value) {
    std::cout << "Inserting " << value << '\n';
    tree.insert(value);
    tree.print(std::cout);
}

template <typename T>
void printElems (AVLTree<T> & tree) {
    bool fst = true;
    tree.each([&fst] (T & val) {
        if (fst)
            fst = false;
        else
            std::cout << ", ";
        std::cout << val;
    });
    std::cout << '\n';
}


int main(int argc, const char * argv[]) {
    
    AVLTree<int> emptyTree {};
    emptyTree.print(std::cout);
    printHasMember(emptyTree, 1);

    AVLTree<int> tree1 { 1 };
    tree1.print(std::cout);
    printHasMember(tree1, 1);
    printHasMember(tree1, 5);

    AVLTree<int> tree2 { 1, 2, 3 };
    tree2.print(std::cout);
    printHasMember(tree2, 1);
    printHasMember(tree2, 2);
    printHasMember(tree2, 3);
    printHasMember(tree2, 4);
    
    AVLTree<int> tree3 { 5, 2, 6, 1, 8, 4, 7 };
    tree3.print(std::cout);
    printHasMember(tree3, 8);
    printHasMember(tree3, 9);
    
    AVLTree<int> fwdTree  { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };
    AVLTree<int> backTree { 10, 9, 8, 7, 6, 5, 4, 3, 2, 1 };
    AVLTree<int> rndTree  { 2, 9, 1, 3, 7, 6, 5, 4, 10, 8 };
    
    
    fwdTree.print(std::cout);
    printElems(fwdTree);
    
    backTree.print(std::cout);
    printElems(backTree);
    
    rndTree.print(std::cout);
    printElems(rndTree);
    
    std::cout << "Removal: \n";
    tree2.print(std::cout);
    printRemoval(tree2, 2);
    printRemoval(tree2, 3);
    printRemoval(tree2, 2);
    
    AVLTree<int> cpy { fwdTree };
    
    fwdTree.print(std::cout);
    printRemoval(fwdTree, 2);
    printElems(fwdTree);
    printRemoval(fwdTree, 1);
    printElems(fwdTree);
    printRemoval(fwdTree, 3);
    printElems(fwdTree);
    printRemoval(fwdTree, 4);
    printInsertion(fwdTree, 4);
    printInsertion(fwdTree, 3);
    printInsertion(fwdTree, 2);
    printInsertion(fwdTree, 1);
    printElems(fwdTree);
    
    std::cout << "Original:\n";
    cpy.print(std::cout);
    printElems(cpy);
}