Saf/Algo/Struct/RedBlackTree.h

Go to the documentation of this file.

/*
    This file is part of Simple Application Framework (Saf) library.
    Copyright (C) 2010 - 2012 Ondrej Danek <ondrej.danek@gmail.com>

    This library is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published 
    by the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    Saf is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Simple Application Framework library. If not, 
    see <http://www.gnu.org/licenses/>.
*/

#ifndef SAF_ALGO_STRUCT_REDBLACKTREE_H
#define SAF_ALGO_STRUCT_REDBLACKTREE_H

#include "../../Type.h"
#include "../../Mem/Alloc.h"
#include "../../Collection/Pair.h"
#include "../Swap.h"
#include "../../Type/IteratorCategory.h"

namespace Saf
{
    namespace Algo
    {        
        namespace Struct
        {
            template <class KeyType, class RecType, class KeySelectorType, class CompType> 
            class RedBlackTree
            {
            protected:
                typedef RedBlackTree<KeyType,RecType,KeySelectorType,CompType> MyType;

            protected:
                class Node
                {
                public:
                    enum Flags
                    {
                        BlackFlag = 0,
                        RedFlag = 1,
                        HeadFlag = 2
                    };

                public:
                    Node *m_parent;
                    Node *m_child[2];
                    RecType m_val;
                    Uint8 m_flags;

                public:
                    Node(const RecType& val, Uint8 flags)
                        : m_val(val), m_flags(flags)
                    {}

                    bool IsHead() const
                    {
                        return (m_flags & HeadFlag) != 0;
                    }

                    bool IsRed() const
                    {
                        return (m_flags & Node::RedFlag) != 0;
                    }

                    Uint8 Colour() const
                    {
                        return m_flags & 1;
                    }

                    void Recolour(Uint8 col)
                    {
                        m_flags = (m_flags & (~1)) | col;
                    }
                };

            public:
                class ConstIterator
                {
                public:
                    typedef RecType ValType;
                    typedef Type::IteratorCategory::BidirectionalCategory CategoryType;

                protected:
                    Node *m_node;

                    // The tree can access the protected members
                    friend class RedBlackTree<KeyType,RecType,KeySelectorType,CompType>;

                public:
                    ConstIterator()
                    {}

                    explicit ConstIterator(Node *node)
                        : m_node(node)
                    {}

                    ConstIterator(const ConstIterator& iter)
                        : m_node(iter.m_node)
                    {}
                    
                    ConstIterator& operator=(const ConstIterator& iter)
                    {
                        m_node = iter.m_node;
                        return *this;
                    }

                    const RecType& operator*() const
                    {
                        return m_node->m_val;
                    }

                    const RecType* operator->() const
                    {
                        return &m_node->m_val;
                    }

                    ConstIterator& operator++()
                    {
                        m_node = TreeNext(m_node);
                        return *this;
                    }

                    ConstIterator operator++(int)
                    {
                        ConstIterator ci(*this);
                        ++(*this);
                        return ci;
                    }

                    ConstIterator& operator--()
                    {
                        m_node = TreePrev(m_node);
                        return *this;
                    }

                    ConstIterator operator--(int)
                    {
                        ConstIterator ci(*this);
                        --(*this);
                        return ci;
                    }

                    bool operator==(const ConstIterator& ci) const
                    {
                        return (m_node == ci.m_node);
                    }

                    bool operator!=(const ConstIterator& ci) const
                    {
                        return (m_node != ci.m_node);
                    }
                };

                class Iterator
                    : public ConstIterator
                {
                public:
                    Iterator()
                        : ConstIterator()
                    {}

                    explicit Iterator(Node *node)
                        : ConstIterator(node)
                    {}

                    Iterator(const Iterator& iter)
                        : ConstIterator(iter)
                    {}
                    
                    Iterator& operator=(const Iterator& iter)
                    {
                        ConstIterator::operator=(iter);
                        return *this;
                    }

                    RecType& operator*() const
                    {
                        return this->m_node->m_val;
                    }

                    RecType* operator->() const
                    {
                        return &this->m_node->m_val;
                    }

                    Iterator& operator++()
                    {
                        ++(*((ConstIterator *)this));
                        return *this;
                    }

                    Iterator operator++(int)
                    {
                        Iterator it(*this);
                        ++(*this);
                        return it;
                    }

                    Iterator& operator--()
                    {
                        --(*((ConstIterator *)this));
                        return *this;
                    }

                    Iterator operator--(int)
                    {
                        Iterator it(*this);
                        --(*this);
                        return it;
                    }
                };

            private:
                Node *m_head;
                Size m_elems;
                CompType m_comp;
                KeySelectorType m_keyExtract;

            protected:
                Node *AllocNode(const RecType& val, Node *parent, Uint8 flags)
                {
                    Node *n = SAF_NEW Node(val, flags);
                    n->m_parent = parent;
                    n->m_child[0] = Head();
                    n->m_child[1] = Head();
                    return n;
                }

                Node* Head() const
                {
                    return m_head;
                }

                Node*& Root() const
                {
                    return Head()->m_parent;
                }

                Node*& Leftmost() const
                {
                    return Head()->m_child[0];
                }

                Node*& Rightmost() const
                {
                    return Head()->m_child[1];
                }

                bool IsLeftChild(Node *n) const
                {
                    return n == n->m_parent->m_child[0];
                }

                void Rotate(Node* root, Int32 dir)
                {
                    Node *grandpa = root->m_parent;
                    Node *nodeA = root->m_child[1-dir];
                    Node *nodeB = nodeA->m_child[dir];

                    nodeA->m_parent = grandpa;
                    if (grandpa->IsHead())
                    {
                        Root() = nodeA;
                    }
                    else
                    {
                        grandpa->m_child[(root == grandpa->m_child[0]) ? 0 : 1] = nodeA;
                    }
                    
                    nodeA->m_child[dir] = root;
                    root->m_parent = nodeA;

                    root->m_child[1-dir] = nodeB;
                    if (!nodeB->IsHead())
                    {
                        nodeB->m_parent = root;
                    }
                }

                Node* TreeLowerBound(const KeyType& key) const
                {
                    Node *node = Root();
                    Node *lowerBound = Head();

                    while (!node->IsHead())
                    {
                        if (m_comp(m_keyExtract(node->m_val), key))
                        {
                            node = node->m_child[1];
                        }
                        else
                        {
                            lowerBound = node;
                            node = node->m_child[0];
                        }
                    }

                    return lowerBound;
                }

                Node* TreeUpperBound(const KeyType& key) const
                {
                    Node *node = Root();
                    Node *upperBound = Head();

                    while (!node->IsHead())
                    {
                        if (m_comp(key, m_keyExtract(node->m_val)))
                        {
                            upperBound = node;
                            node = node->m_child[0];
                        }
                        else
                        {
                            node = node->m_child[1];
                        }
                    }

                    return upperBound;
                }

                static Node* TreePrev(Node *node)
                {
                    if (node->IsHead())
                    {
                        return node->m_child[1];
                    }

                    if (!node->m_child[0]->IsHead())
                    {
                        return TreeMax(node->m_child[0]);
                    }

                    Node *parent = node->m_parent;
                    while (!parent->IsHead() && node == parent->m_child[0])
                    {
                        node = parent;
                        parent = parent->m_parent;
                    }

                    return parent;
                }

                static Node* TreeNext(Node *node)
                {
                    if (node->IsHead())
                    {
                        return node;
                    }

                    if (!node->m_child[1]->IsHead())
                    {
                        return TreeMin(node->m_child[1]);
                    }

                    Node *parent = node->m_parent;
                    while (!parent->IsHead() && node == parent->m_child[1])
                    {
                        node = parent;
                        parent = parent->m_parent;
                    }

                    return parent;
                }

                static Node* TreeMin(Node *root)
                {
                    while (!root->m_child[0]->IsHead())
                    {
                        root = root->m_child[0];
                    }

                    return root;
                }

                static Node* TreeMax(Node *root)
                {
                    while (!root->m_child[1]->IsHead())
                    {
                        root = root->m_child[1];
                    }

                    return root;
                }

            private:
                void Init()
                {
                    // Allocate a head node, but do not construct the RecType object                    
                    Mem::OperatorNew(m_head, SAF_SOURCE_LOCATION);
                    m_head->m_flags = Node::HeadFlag | Node::BlackFlag;
                    
                    // Init tree
                    Root() = Head();
                    Leftmost() = Head();
                    Rightmost() = Head();
                }

            public:
                RedBlackTree() 
                    : m_head(nullptr), m_elems(0)
                {
                    Init();
                }

                RedBlackTree(const MyType &rbt) 
                    : m_head(nullptr), m_elems(0)
                { 
                    Init();
                    *this = rbt; 
                }

                ~RedBlackTree() 
                { 
                    Free();
                    Mem::OperatorDelete(m_head);
                }

                MyType& Free() 
                {                    
                    if (!IsEmpty())
                    {
                        Node *node = Root(), *nextNode;

                        while (node != Head())
                        {
                            if (!node->m_child[0]->IsHead())
                            {
                                nextNode = node->m_child[0];
                                node->m_child[0] = Head();
                                node = nextNode;
                            }
                            else if (!node->m_child[1]->IsHead())
                            {
                                nextNode = node->m_child[1];
                                node->m_child[1] = Head();
                                node = nextNode;
                            }
                            else
                            {
                                nextNode = node->m_parent;
                                Mem::Delete(node);
                                node = nextNode;
                            }
                        }

                        Root() = Head();
                        Leftmost() = Head();
                        Rightmost() = Head();
                        m_elems = 0;
                    }
                
                    return *this;
                }

                MyType& operator=(const MyType& rbt)
                {
                    if (&rbt != this)
                    {
                        Free();

                        if (!rbt.IsEmpty())
                        {
                            try
                            {
                                Root() = AllocNode(rbt.Root()->m_val, Head(), Node::BlackFlag);
                                m_elems = 1;

                                Node *nDst = Root(), *nSrc = rbt.Root();

                                while (nDst != Head())
                                {
                                    if (nDst->m_child[0]->IsHead() && !nSrc->m_child[0]->IsHead())
                                    {
                                        nSrc = nSrc->m_child[0];
                                        nDst->m_child[0] = AllocNode(nSrc->m_val, nDst, nSrc->m_flags);
                                        nDst = nDst->m_child[0];
                                        m_elems++;                                
                                    }
                                    else if (nDst->m_child[1]->IsHead() && !nSrc->m_child[1]->IsHead())
                                    {
                                        nSrc = nSrc->m_child[1];
                                        nDst->m_child[1] = AllocNode(nSrc->m_val, nDst, nSrc->m_flags);
                                        nDst = nDst->m_child[1];
                                        m_elems++;                                
                                    }
                                    else
                                    {
                                        nDst = nDst->m_parent;
                                        nSrc = nSrc->m_parent;
                                    }
                                }

                                Leftmost() = TreeMin(Root());
                                Rightmost() = TreeMax(Root());
                            }
                            catch (...)
                            {
                                Free();
                                throw;
                            }
                        }
                    }

                    return *this;
                }

                MyType& Swap(MyType& rbt)
                {
                    if (this != &rbt)
                    {
                        Algo::Swap(m_head, rbt.m_head);
                        Algo::Swap(m_elems, rbt.m_elems);
                        Algo::Swap(m_comp, rbt.m_comp);
                        Algo::Swap(m_keyExtract, rbt.m_keyExtract);
                    }

                    return *this;
                }

                bool IsEmpty() const 
                { 
                    return (m_elems == 0);
                }

                Size Elements() const
                {
                    return m_elems;
                }

                Collection::Pair<Iterator,bool> Insert(const RecType& val)
                {
                    // Handle inserts into an empty tree
                    if (IsEmpty())
                    {                        
                        Node *newRoot = AllocNode(val, Head(), Node::BlackFlag);
                        
                        Root() = newRoot;
                        Leftmost() = newRoot;
                        Rightmost() = newRoot;
                        m_elems = 1;

                        return Collection::Pair<Iterator,bool>(Iterator(newRoot), true);
                    }

                    // Non-empty tree -> find where to put the new node
                    Node *node = Root(), *parent = Head(), *lowerBound = Head();
                    Int32 dir = 0;

                    while (!node->IsHead())
                    {
                        if (m_comp(m_keyExtract(val), m_keyExtract(node->m_val)))
                        {
                            dir = 0;
                        }
                        else
                        {
                            dir = 1;
                            lowerBound = node;
                        }

                        parent = node;
                        node = node->m_child[dir];
                    }                    

                    // Check if unique by checking the value of the lower bound
                    if (!lowerBound->IsHead())
                    {
                        if (!m_comp(m_keyExtract(lowerBound->m_val), m_keyExtract(val)))
                        {
                            // Key with equivalent ordering already exists in the tree
                            return Collection::Pair<Iterator,bool>(Iterator(lowerBound), false);
                        }
                    }

                    // Insert a new element
                    Node *newNode = AllocNode(val, parent, Node::RedFlag);
                    m_elems++;

                    node = newNode;
                    parent->m_child[dir] = node;

                    if (parent == Leftmost() && dir == 0)
                    {
                        Leftmost() = node; // Update the left-most element
                    }
                    else if (parent == Rightmost() && dir != 0)
                    {
                        Rightmost() = node; // Update the right-most element
                    }

                    // Rebalance
                    while (parent->IsRed())
                    {
                        dir = IsLeftChild(parent) ? 0 : 1;
                        
                        Node *grandpa = parent->m_parent;
                        Node *uncle = grandpa->m_child[1-dir];

                        if (uncle->IsRed())
                        {
                            // Uncle is red -> move red to grandpa
                            parent->Recolour(Node::BlackFlag);
                            uncle->Recolour(Node::BlackFlag);
                            grandpa->Recolour(Node::RedFlag);
                            node = grandpa;
                        }
                        else
                        {
                            // Uncle is black -> rotate subtrees
                            if (node == parent->m_child[1-dir])
                            {
                                Rotate(parent, dir);
                                node = parent;
                                parent = node->m_parent;
                            }

                            parent->Recolour(Node::BlackFlag);
                            grandpa->Recolour(Node::RedFlag);
                            Rotate(grandpa, 1-dir);
                        }

                        parent = node->m_parent;
                    }

                    Root()->Recolour(Node::BlackFlag); // Root is always black
                    return Collection::Pair<Iterator,bool>(Iterator(newNode), true);
                }

                Iterator Erase(Iterator iter)
                {
                    // Iterator to be returned
                    Iterator succIter(iter.m_node);
                    ++succIter;

                    // Erase
                    Node *node = iter.m_node;
                    Node *left = node->m_child[0], *right = node->m_child[1];
                    Node *parent = node->m_parent;
                    bool needRebalance;

                    // Decide which element will replace the erased node and relink
                    if (left->IsHead() || right->IsHead())
                    {                        
                        // Lift left or right child
                        Node *next = left->IsHead() ? right : left;
                        needRebalance = !node->IsRed();

                        // Relink the subtree
                        if (!next->IsHead())
                        {
                            next->m_parent = parent;
                        }

                        if (parent->IsHead())
                        {
                            Root() = next;
                        }
                        else
                        {
                            parent->m_child[(node == parent->m_child[0]) ? 0 : 1] = next;
                        }

                        // Fix left-most and right-most elements
                        if (node == Leftmost())
                        {
                            Leftmost() = succIter.m_node;
                        }

                        if (node == Rightmost())
                        {
                            Rightmost() = left->IsHead() ? parent : TreeMax(left);
                        }

                        // Proceed with rebalancing next                        
                        node = next;
                    }
                    else
                    {
                        // Lift successor
                        Node *next = succIter.m_node, *nextParent = next->m_parent;                        

                        // Recolour
                        needRebalance = !next->IsRed();
                        next->Recolour(node->Colour());

                        // The new root of the subtree is next
                        next->m_parent = parent;
                        next->m_child[0] = left;
                        left->m_parent = next;

                        if (parent->IsHead())
                        {
                            Root() = next;
                        }
                        else
                        {
                            parent->m_child[(node == parent->m_child[0]) ? 0 : 1] = next;
                        }

                        // Right child of the succesor will be the first rebalanced node
                        node = next->m_child[1];                        

                        // Fix the right half of the subtree
                        if (next == right)
                        {
                            // Successor is next to the erased node
                            parent = next;
                        }
                        else
                        {
                            parent = nextParent;
                            if (!node->IsHead())
                            {
                                node->m_parent = parent;
                            }
                            parent->m_child[0] = node;
                            next->m_child[1] = right;
                            right->m_parent = next;
                        }
                    }

                    // Rebalance
                    if (needRebalance)
                    {
                        while (!node->IsRed() && node != Root())
                        {
                            Int32 dir = (node == parent->m_child[0]) ? 0 : 1; // Fix left or right subtree
                            Node *sibling = parent->m_child[1-dir];

                            if (sibling->IsRed())
                            {
                                // Rotate red from left or right subtree
                                sibling->Recolour(Node::BlackFlag);
                                parent->Recolour(Node::RedFlag);
                                Rotate(parent, dir);
                                sibling = parent->m_child[1-dir];
                            }

                            if (!sibling->m_child[0]->IsRed() && !sibling->m_child[1]->IsRed())
                            {
                                // Redden left/right subtree with black children
                                sibling->Recolour(Node::RedFlag);
                                node = parent;
                                parent = node->m_parent;
                            }
                            else
                            {
                                // Rearrange left/right subtree
                                if (!sibling->m_child[1-dir]->IsRed())
                                {
                                    sibling->m_child[dir]->Recolour(Node::BlackFlag);
                                    sibling->Recolour(Node::RedFlag);
                                    Rotate(sibling, 1-dir);
                                    sibling = parent->m_child[1-dir];
                                }

                                sibling->Recolour(parent->Colour());
                                parent->Recolour(Node::BlackFlag);
                                sibling->m_child[1-dir]->Recolour(Node::BlackFlag);
                                Rotate(parent, dir);
                                break;
                            }
                        }

                        // Stopping node is black
                        node->Recolour(Node::BlackFlag);
                    }

                    // Free node
                    --m_elems;
                    Mem::Delete(iter.m_node);                    
                    
                    return succIter;
                }

                Iterator LowerBound(const KeyType& key)
                {
                    return Iterator(TreeLowerBound(key));
                }

                ConstIterator LowerBound(const KeyType& key) const
                {
                    return ConstIterator(TreeLowerBound(key));
                }

                Iterator UpperBound(const KeyType& key)
                {
                    return Iterator(TreeUpperBound(key));
                }

                ConstIterator UpperBound(const KeyType& key) const
                {
                    return ConstIterator(TreeUpperBound(key));
                }

                Iterator Find(const KeyType& key)
                {
                    Iterator it = LowerBound(key);

                    if (it == End() || m_comp(key, m_keyExtract(*it)))
                    {
                        return End();
                    }

                    return it;
                }

                ConstIterator Find(const KeyType& key) const
                {
                    ConstIterator it = LowerBound(key);

                    if (it == End() || m_comp(key, m_keyExtract(*it)))
                    {
                        return End();
                    }

                    return it;
                }

                bool Contains(const KeyType& key) const
                {
                    return (Find(key) != End());
                }

                Iterator Begin()
                {
                    return Iterator(Leftmost());
                }

                ConstIterator Begin() const
                {
                    return ConstIterator(Leftmost());
                }

                Iterator End()
                {
                    return Iterator(Head());
                }

                ConstIterator End() const
                {
                    return ConstIterator(Head());
                }

                Iterator Front()
                {
                    return Begin();
                }

                ConstIterator Front() const
                {
                    return Begin();
                }

                Iterator Back()
                {
                    return Iterator(Rightmost());
                }

                ConstIterator Back() const
                {
                    return ConstIterator(Rightmost());
                }
            };
        }

        // Swap specialization for red/black trees.
        template <class KeyType, class RecType, class KeySelectorType, class CompType> 
        struct SwapFunc< Struct::RedBlackTree<KeyType,RecType,KeySelectorType,CompType> >
        {
            void operator()(Struct::RedBlackTree<KeyType,RecType,KeySelectorType,CompType>& x, 
                Struct::RedBlackTree<KeyType,RecType,KeySelectorType,CompType>& y)
            {
                x.Swap(y);
            }
        };
    }
}

#endif