Initial commit

BlueNoise/WeightedDiscretePDF.cpp

@@ -0,0 +1,313 @@
+// $Id: WeightedDiscretePDF.cpp,v 1.4 2006/07/07 05:54:31 zr Exp $
+
+#include <stdexcept>
+#include "WeightedDiscretePDF.h"
+
+//
+
+template <class T>
+WDPDF_Node<T>::WDPDF_Node(T key_, float weight_, WDPDF_Node<T> *parent_)
+{
+	m_mark = false;
+
+	key = key_;
+	weight = weight_;
+	sumWeights = 0;
+	left = right = 0;
+	parent = parent_;
+}
+
+template <class T>
+WDPDF_Node<T>::~WDPDF_Node()
+{
+	if (left) delete left;
+	if (right) delete right;
+}
+
+//
+
+template <class T>
+WeightedDiscretePDF<T>::WeightedDiscretePDF()
+{
+	m_root = 0;
+}
+
+template <class T>
+WeightedDiscretePDF<T>::~WeightedDiscretePDF()
+{
+	if (m_root) delete m_root;
+}
+
+template <class T>
+void WeightedDiscretePDF<T>::insert(T item, float weight)
+{
+	WDPDF_Node<T> *p=0, *n=m_root;
+
+	while (n) {
+		if (n->leftIsRed() && n->rightIsRed())
+			split(n);
+
+		p = n;
+		if (n->key==item) {
+			throw std::domain_error("insert: argument(item) already in tree");
+		} else {
+			n = (item<n->key)?n->left:n->right;
+		}
+	}
+
+	n = new WDPDF_Node<T>(item, weight, p);
+
+	if (!p) {
+		m_root = n;
+	} else {
+		if (item<p->key) {
+			p->left = n;
+		} else {
+			p->right = n;
+		}
+
+		split(n);
+	}
+
+	propogateSumsUp(n);
+}
+
+template <class T>
+void WeightedDiscretePDF<T>::remove(T item)
+{
+	WDPDF_Node<T> **np = lookup(item, 0);
+	WDPDF_Node<T> *child, *n = *np;
+
+	if (!n) {
+		throw std::domain_error("remove: argument(item) not in tree");
+	} else {
+		if (n->left) {
+			WDPDF_Node<T> **leftMaxp = &n->left;
+
+			while ((*leftMaxp)->right)
+				leftMaxp = &(*leftMaxp)->right;
+
+			n->key = (*leftMaxp)->key;
+			n->weight = (*leftMaxp)->weight;
+
+			np = leftMaxp;
+			n = *np;
+		}
+
+			// node now has at most one child
+
+		child = n->left?n->left:n->right;
+		*np = child;
+
+		if (child) {
+			child->parent = n->parent;
+
+			if (n->isBlack()) {
+				lengthen(child);
+			}
+		}
+
+		propogateSumsUp(n->parent);
+
+		n->left = n->right = 0;
+		delete n;
+	}
+}
+
+template <class T>
+void WeightedDiscretePDF<T>::update(T item, float weight)
+{
+	WDPDF_Node<T> *n = *lookup(item, 0);
+
+	if (!n) {
+		throw std::domain_error("update: argument(item) not in tree");
+	} else {
+		float delta = weight - n->weight;
+		n->weight = weight;
+
+		for (; n; n=n->parent) {
+			n->sumWeights += delta;
+		}
+	}
+}
+
+template <class T>
+T WeightedDiscretePDF<T>::choose(float p)
+{
+	if (p<0.0 || p>=1.0) {
+		throw std::domain_error("choose: argument(p) outside valid range");
+	} else if (!m_root) {
+		throw std::logic_error("choose: choose() called on empty tree");
+	} else {
+		float w = m_root->sumWeights * p;
+		WDPDF_Node<T> *n = m_root;
+
+		while (1) {
+			if (n->left) {
+				if (w<n->left->sumWeights) {
+					n = n->left;
+					continue;
+				} else {
+					w -= n->left->sumWeights;
+				}
+			}
+			if (w<n->weight || !n->right) {
+				break; // !n->right condition shouldn't be necessary, just sanity check
+			}
+			w -= n->weight;
+			n = n->right;
+		}
+
+		return n->key;
+	}
+}
+
+template <class T>
+bool WeightedDiscretePDF<T>::inTree(T item)
+{
+	WDPDF_Node<T> *n = *lookup(item, 0);
+
+	return !!n;
+}
+
+//
+
+template <class T>
+WDPDF_Node<T> **WeightedDiscretePDF<T>::lookup(T item, WDPDF_Node<T> **parent_out)
+{
+	WDPDF_Node<T> *n, *p=0, **np=&m_root;
+
+	while ((n = *np)) {
+		if (n->key==item) {
+			break;
+		} else {
+			p = n;
+			if (item<n->key) {
+				np = &n->left;
+			} else {
+				np = &n->right;
+			}
+		}
+	}
+
+	if (parent_out)
+		*parent_out = p;
+	return np;
+}
+
+template <class T>
+void WeightedDiscretePDF<T>::split(WDPDF_Node<T> *n)
+{
+	if (n->left) n->left->markBlack();
+	if (n->right) n->right->markBlack();
+
+	if (n->parent) {
+		WDPDF_Node<T> *p = n->parent;
+
+		n->markRed();
+
+		if (p->isRed()) {
+			p->parent->markRed();
+
+				// not same direction
+			if (!(	(n==p->left && p==p->parent->left) || 
+					(n==p->right && p==p->parent->right))) {
+			  rotate(n);
+			  p = n;
+			}
+
+			rotate(p);
+			p->markBlack();
+		}
+	}
+}
+
+template <class T>
+void WeightedDiscretePDF<T>::rotate(WDPDF_Node<T> *n)
+{
+	WDPDF_Node<T> *p=n->parent, *pp=p->parent;
+
+	n->parent = pp;
+	p->parent = n;
+
+	if (n==p->left) {
+		p->left = n->right;
+		n->right = p;
+		if (p->left) p->left->parent = p;
+	} else {
+		p->right = n->left;
+		n->left = p;
+		if (p->right) p->right->parent = p;
+	}
+
+	n->setSum();
+	p->setSum();
+	
+	if (!pp) {
+		m_root = n;
+	} else {
+		if (p==pp->left) {
+			pp->left = n;
+		} else {
+			pp->right = n;
+		}
+	}
+}
+
+template <class T>
+void WeightedDiscretePDF<T>::lengthen(WDPDF_Node<T> *n)
+{
+	if (n->isRed()) {
+		n->markBlack();
+	} else if (n->parent) {
+		WDPDF_Node<T> *sibling = n->sibling();
+
+		if (sibling && sibling->isRed()) {
+			n->parent->markRed();
+			sibling->markBlack();
+
+			rotate(sibling); // node sibling is now old sibling child, must be black
+			sibling = n->sibling();
+		}
+
+		// sibling is black
+
+		if (!sibling) {
+			lengthen(n->parent);
+		} else if (sibling->leftIsBlack() && sibling->rightIsBlack()) {
+			if (n->parent->isBlack()) {
+				sibling->markRed();
+				lengthen(n->parent);
+			} else {
+				sibling->markRed();
+				n->parent->markBlack();
+			}
+		} else {
+			if (n==n->parent->left && sibling->rightIsBlack()) {
+				rotate(sibling->left); // sibling->left must be red
+				sibling->markRed();
+				sibling->parent->markBlack();
+				sibling = sibling->parent;
+			} else if (n==n->parent->right && sibling->leftIsBlack()) {
+				rotate(sibling->right); // sibling->right must be red
+				sibling->markRed();
+				sibling->parent->markBlack();
+				sibling = sibling->parent;
+			}
+
+			// sibling is black, and sibling's far child is red
+
+			rotate(sibling);
+			if (n->parent->isRed()) sibling->markRed();
+			sibling->left->markBlack();
+			sibling->right->markBlack();
+		}
+	}
+}
+
+template <class T>
+void WeightedDiscretePDF<T>::propogateSumsUp(WDPDF_Node<T> *n)
+{
+	for (; n; n=n->parent)
+		n->setSum();
+}