57 lines
1.4 KiB
C++
57 lines
1.4 KiB
C++
// $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);
|
|
};
|