WeightedDiscretePDF.h

// $Id: WeightedDiscretePDF.h,v 1.4 2006/07/07 05:54:31 zr Exp $

template <class T>
class WDPDF_Node
{
private:
        bool m_mark;

public:
        WDPDF_Node<T> *parent, *left, *right;
        T key;
        float weight, sumWeights;

public:
        WDPDF_Node(T key_, float weight_, WDPDF_Node<T> *parent_);
        ~WDPDF_Node();

        WDPDF_Node<T> *sibling() { return this==parent->left?parent->right:parent->left; }

        void markRed() { m_mark = true; }
        void markBlack() { m_mark = false; }
        bool isRed() { return m_mark; }
        bool isBlack() { return !m_mark; }
        bool leftIsBlack() { return !left || left->isBlack(); }
        bool rightIsBlack() { return !right || right->isBlack(); }
        bool leftIsRed() { return !leftIsBlack(); }
        bool rightIsRed() { return !rightIsBlack(); }
        void setSum() { sumWeights = weight + (left?left->sumWeights:0) + (right?right->sumWeights:0); }
};

template <class T>
class WeightedDiscretePDF
{
private:
        WDPDF_Node<T> *m_root;

public:
        WeightedDiscretePDF();
        ~WeightedDiscretePDF();

        void insert(T item, float weight);
        void update(T item, float newWeight);
        void remove(T item);
        bool inTree(T item);
        
                /* pick a tree element according to its
                 * weight. p should be in [0,1).
                 */
        T choose(float p);

private:
        WDPDF_Node<T> **lookup(T item, WDPDF_Node<T> **parent_out);
        void split(WDPDF_Node<T> *node);
        void rotate(WDPDF_Node<T> *node);
        void lengthen(WDPDF_Node<T> *node);
        void propogateSumsUp(WDPDF_Node<T> *n);
};

#include "WeightedDiscretePDF.cpp"

---