/*-
 * Copyright (c) 2000 Qube Software, Ltd.
 * All rights reserved.
 * 
 * 	$Id: wtree.h,v 1.3 2000/05/26 11:06:01 dfr Exp $
 */

#ifndef __WTree_h__
#define __WTree_h__

#include        <iostream>
#include        <assert.h>
#include        "vector.h"
#include        "math.h"

// #define TREE_DEBUG

/* Forward Decls */
template<class K> struct WTree;

template<class K> struct WTreeNode
{
    typedef WTreeNode<K>            nodeType;

    friend struct WTree<K>;

    // Constraints on K
    // int operator==(const K&) const;
    // int operator<=(const k&) const;
    // K::K();

                                WTreeNode(const K& key)
                                    : key_(key) { init(); }

    // Destructor
                                ~WTreeNode();

    // Accessors
    const K&                    key() const { return key_; }
    nodeType*                   left() const { return left_; }
    nodeType*                   right() const { return right_; }

    // Features
    unsigned int                depth() const;

private:

    void                        init();

    K                           key_;
    unsigned int                weight_;
    nodeType*                   left_;
    nodeType*                   right_;
};


template<class K> struct WTree
{
    typedef WTreeNode<K>       nodeType;


    // Constructors
                                WTree() { top_ = 0; }

    // Destructor
                                ~WTree() { empty(); }

    // Accessors
    nodeType*                   top() const { return top_; }

    // Features
    const K*                    min() const;
    const K*                    max() const;
    bool                        insert(const K& key);
    bool                        remove(const K&);
    unsigned int                remove(const K&, const K&);
    unsigned int                removeC(const K&, const K&);
    void                        empty() { delete top_; top_ = 0; }
    const K*                    find(const K& key) const;
    bool                        find(const K& key, K*&, K*&) const;
    unsigned int                depth() const 
                                        { return top_ ? top_->depth() : 0; }
    void                        checkTree() const;
    static int                  checkDepth(nodeType*);

private:

    static void                 adjust(nodeType*&);
    static bool                 insert(const K& key, nodeType*&);
    static void                 lrot(nodeType*&);
    static void                 rrot(nodeType*&);
    static bool                 remove(const K& key, nodeType*&);
    static void                 checkrots(nodeType*&);
    static unsigned int         remove(const K&, const K&, 
                                       const K&, const K&, nodeType*&);
    static unsigned int         removeC(const K&, const K&, 
                                     const K&, const K&, nodeType*&);
    static unsigned int         prune(nodeType*&);
    static unsigned int         wt(const nodeType* t)
                                        { return t ? t->weight_ :  1; }

    nodeType*                   top_;
};

/*********** WTreeNode<K>, methods  **************************/

template <class K> void WTreeNode<K>::init()
{
    left_ = 0;
    right_ = 0;
    weight_ = 0;
}

template<class K> WTreeNode<K>::~WTreeNode()
{
    delete left_;
    delete right_;
}

template<class K> unsigned int WTreeNode<K>::depth() const
{
    unsigned int dl, dr;
    dl = left_ ? left_->depth() : 0;
    dr = right_ ? right_->depth() : 0;
    if (dl > dr) dr = dl;
    return dr + 1;
}

/*********** WTree<K>, methods  ********************************/

template <class K>
bool WTree<K>::insert(const K& key)
{
    bool v = insert(key, top_);

#ifdef TREE_DEBUG
    checkTree();
#endif
    return v;
}

template <class K>
void WTree<K>::adjust(nodeType*& node)
{
    nodeType* t = node;

    unsigned int a = wt(t->left_);
    unsigned int b = wt(t->right_);
        
    t->weight_ = a + b;
    if (29*a > 70*b) {
        b = wt(t->left_->left_);
        if (169*b < 70*a) lrot(t->left_);
        rrot(node);
    }
    else if (239*a < 99*b) {
        a = wt(t->right_->left_);
        if (169*a > 99*b) rrot(t->right_);
        lrot(node);
    }
}

template <class K>
bool WTree<K>::insert(const K& key, nodeType*& node)
{
    nodeType* t = node;
    if (!t) {
        node = new WTreeNode<K>(key);
        node->weight_ = 2;
    }
    else if (t->key_ == key) {
        /* install new version */
        t->key_ = key;  
        return false;           // already present.
    }
    else {
        if (t->key_ < key) insert(key, t->right_);
        else insert(key, t->left_);
        adjust(node);

    }
    return true;
}

template<class K> void WTree<K>::lrot(nodeType*& node)
{
    nodeType* t = node;
    nodeType* x = t->right_;
    node = x;
    t->right_ = x->left_;
    x->left_ = t;
    
    t->weight_ = wt(t->left_) + wt(t->right_);
    x->weight_ = t->weight_ + wt(x->right_);
}

template<class K> void WTree<K>::rrot(nodeType*& node)
{
    nodeType* t = node;
    nodeType* x = t->left_;
    node = x;
    t->left_ = x->right_;
    x->right_ = t;
    
    t->weight_ = wt(t->left_) + wt(t->right_);
    x->weight_ = t->weight_ + wt(x->left_);
}

template<class K> bool WTree<K>::remove(const K& key)
{
    bool v = remove(key, top_);
#ifdef TREE_DEBUG
    checkTree();
#endif
    return v;
}

template<class K> unsigned int
WTree<K>::remove(const K& keyMin, const K& keyMax)
{
    if (!top_) return 0;

    unsigned int n =
        remove(keyMin, keyMax, *min(), *max(), top_);
#ifdef TREE_DEBUG
    checkTree();
#endif
    return n;
}

template<class K> unsigned int
WTree<K>::removeC(const K& keyMin, const K& keyMax)
{
    if (!top_) return 0;

    unsigned int n =
        removeC(keyMin, keyMax, *min(), *max(), top_);
#ifdef TREE_DEBUG
    checkTree();
#endif
    return n;
}

template<class K> bool WTree<K>::remove(const K& key, nodeType*& node)
{
    nodeType* t = node;
    if (!t) return false;

    bool result;
    if (t->key_ < key) result = remove(key, t->right_);
    else if (key < t->key_) result = remove(key, t->left_);
    else if (!t->left_) {
        node = t->right_;
        t->right_ = 0;
        delete t;
        result = true;
    }
    else if (!t->right_) {
        node = t->left_;
        t->left_ = 0;
        delete t;
        result = true;
    }
    else {
        nodeType** np = &t->left_;
        nodeType* n = *np;
        while (n->right_) {
            np = &n->right_;
            n = *np;
        }

        node = n;
        *np = t;

        nodeType* x = n->left_;
        n->right_ = t->right_;
        n->left_ = t->left_;
        t->right_ = 0;
        t->left_ = x;
        unsigned int w = t->weight_;
        t->weight_ = n->weight_;
        n->weight_ = w;

        result = remove(key, n->left_);
    }
    if (node) adjust(node);

    return result;
}

template<class K> unsigned int
WTree<K>::remove(const K& keyMin, const K& keyMax,
                 const K& inf, const K& sup,
                 nodeType*& node)
{
    /* Remove all in [keyMin, keyMax) */

    unsigned int n = 0;
    nodeType* t = node;
    if (t) {
        bool l = (keyMin <= t->key_) != 0;
        bool h = (t->key_ < keyMax) != 0;
        
        if (!l) {
            n += remove(keyMin, keyMax, t->key_, sup, t->right_);
        }
        else if (!h) {
            n += remove(keyMin, keyMax, inf, t->key_, t->left_);
        }
        else {
            if (keyMin <= inf && sup < keyMax) {
                n += prune(node);
            }
            else {
                n += remove(keyMin, keyMax, inf, t->key_, t->left_);
                n += remove(keyMin, keyMax, t->key_, sup, t->right_);

                /* Combine two trees */
                checkrots(node);

                /* Remove this node */
                remove(t->key_, node);
                ++n;
            }
        }
        if (node) checkrots(node);
    }
    return n;
}

template<class K> unsigned int
WTree<K>::removeC(const K& keyMin, const K& keyMax,
                       const K& inf, const K& sup,
                       nodeType*& node)
{
    /* Remove all in [keyMin, keyMax] */

    unsigned int n = 0;
    nodeType* t = node;
    if (t) {
        bool l = (keyMin <= t->key_) != 0;
        bool h = (t->key_ <= keyMax) != 0;
        
        if (!l) {
            n += removeC(keyMin, keyMax, t->key_, sup, t->right_);
        }
        else if (!h) {
            n += removeC(keyMin, keyMax, inf, t->key_, t->left_);
        }
        else {
            if (keyMin <= inf && sup <= keyMax) {
                n += prune(node);
            }
            else {
                n += removeC(keyMin, keyMax, inf, t->key_, t->left_);
                n += removeC(keyMin, keyMax, t->key_, sup, t->right_);

                /* Combine two trees */
                checkrots(node);

                /* Remove this node */
                remove(t->key_, node);
                ++n;
            }
        }
        if (node) checkrots(node);
    }
    return n;
}

template<class K>
const K* WTree<K>::find(const K& key) const
{
    nodeType* node = top_;
    
    while (node) {
        if (node->key_ == key) return &node->key_;
        if (node->key_ < key) node = node->right_;
        else node = node->left_;
    }
    return 0;
}

template<class K>
bool WTree<K>::find(const K& key, K*& a, K*& b) const
{
    /* Look for `key', but find `a' and `b' such that:
     * a <= key < b or
     * a <= key, b = 0 or       (key >= all tree).
     * a = 0, key < b,          (key less than whole tree).
     * a = 0, b = 0             (tree empty).
     */

    nodeType* t = top_;
    a = 0;
    b = 0;
    
    while (t) {
        if (t->key_ == key) {
            a = &t->key_;
            
            /* sup(key) = b is the least node of t->right. */
            t = t->right_;
            if (t) {
                while (t->left_) t = t->left_;
                b = &t->key_;
            }
            return true;
        }
        else if (t->key_ < key) {
            if (!a || *a < t->key_) a = &t->key_;
            t = t->right_;
        }
        else {
            if (!b || t->key_ < *b) b = &t->key_;
            t = t->left_;
        }
    }
    return false;
}

template <class K>
void WTree<K>::checkrots(nodeType*& node)
{
    /* Balance `node' assuming left and right are already balanced,
     * but both might have changed.
     */
    nodeType* t = node;

    unsigned int a = wt(t->left_);
    unsigned int b = wt(t->right_);

    t->weight_ = a + b;
    
    if (29*a > 70*b) {
        do {
            b = wt(t->left_->left_);
            if (169*b < 70*a) {
                lrot(t->left_);
                rrot(node);
                checkrots(node->left_);
            }
            else rrot(node);

            t = node;
            checkrots(t->right_);
            
            a = wt(t->left_);
            b = wt(t->right_);

        } while (29*a > 70*b);
    }
    else if (239*a < 99*b) {
        do {
            a = wt(t->right_->left_);
            if (169*a > 99*b) {
                rrot(t->right_);
                lrot(node);
                checkrots(node->right_);
            }
            else lrot(node);

            t = node;
            checkrots(t->left_);

            a = wt(t->left_);
            b = wt(t->right_);
        } while (239*a < 99*b);

    }
}

template<class K>
const K* WTree<K>::min() const
{
    nodeType* t = top_;
    if (t) {
        while (t->left_) t = t->left_;
        return &t->key_;
    }
    return 0;
}

template<class K>
const K* WTree<K>::max() const
{
    nodeType* t = top_;
    if (t) {
        while (t->right_) t = t->right_;
        return &t->key_;
    }
    return 0;
}

template<class K>
unsigned int WTree<K>::prune(nodeType*& node)
{
    unsigned int n = 0;
    nodeType* t = node;
    if (t) {
        n += prune(t->left_);
        n += prune(t->right_);

        delete t;
        node = 0;
        ++n;
    }
    return n;
}

#ifdef TREE_DEBUG
template<class K> void WTree<K>::checkTree() const
{
    checkDepth(top_);
}

template<class K> int WTree<K>::checkDepth(nodeType* node)
{
    if (node) {
        int wl = checkDepth(node->left_);
        int wr = checkDepth(node->right_);

        if (node->weight_ != wl + wr) {
            cout << "Node weight " << node->weight_
                 << " not correct " << wl + wr << endl;
            assert(0);
        }
        
        float p = ((float)wl)/(wl + wr);
        float a = 1.0f - sqrtf(0.5f);

        if (p < a || p > 1.0f - a) {
            cout << "Tree out of balance, left = " << wl
                 << " right = " << wr << endl;
            cout << "p(t) = " << p << endl;
            cout << "allowed range: " << a << " to " << 1.0f - a << endl;
            assert(0);
        }

        /* Check key ordering */
        if (node->left_ && !(node->left_->key_ < node->key_)) {
                cout << "order failed (left), " << node->left_->key_
                     << ' ' << node->key_ << endl; 
        }
        if (node->right_ && !(node->key_ < node->right_->key_)) {
                cout << "order failed (right), " << node->key_
                     << ' ' << node->right_->key_ << endl; 
        }

        return wl + wr;
    }
    else {
        return 1;
    }
}

#endif /* TREE_DEBUG */



/***********************************************************************/

template<class K>
std::ostream& operator<<(std::ostream& os, const WTreeNode<K>* node)
{
    os << '(';

    if (node) {
        os << node->key() << ' ';
        os << node->left() << ' ' << node->right();
    }
    os << ')';

    return os;
}

template<class K>
std::ostream& operator<<(std::ostream& os, const WTree<K>& tree)
{
    return os << tree.top();
}

#endif /* __WTree_h__ */


