klee
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Pages
DiscretePDF.inc
Go to the documentation of this file.
1 //===- DiscretePDF.inc - --*- C++ -*-===//
2 
3 //
4 
5 namespace klee {
6 
7 template <class T>
8 class DiscretePDF<T>::Node
9 {
10 private:
11  bool m_mark;
12 
13 public:
14  Node *parent, *left, *right;
15  T key;
16  weight_type weight, sumWeights;
17 
18 public:
19  Node(T key_, weight_type weight_, Node *parent_);
20  ~Node();
21 
22  Node *sibling() { return this==parent->left?parent->right:parent->left; }
23 
24  void markRed() { m_mark = true; }
25  void markBlack() { m_mark = false; }
26  bool isBlack() { return !m_mark; }
27  bool leftIsBlack() { return !left || left->isBlack(); }
28  bool rightIsBlack() { return !right || right->isBlack(); }
29  void setSum() {
30  sumWeights = weight;
31  if (left) sumWeights += left->sumWeights;
32  if (right) sumWeights += right->sumWeights;
33  }
34 };
35 
37 
38 template <class T>
39 DiscretePDF<T>::Node::Node(T key_, weight_type weight_, Node *parent_) {
40  m_mark = false;
41 
42  key = key_;
43  weight = weight_;
44  sumWeights = 0;
45  left = right = 0;
46  parent = parent_;
47 }
48 
49 template <class T>
50 DiscretePDF<T>::Node::~Node() {
51  if (left) delete left;
52  if (right) delete right;
53 }
54 
55 //
56 
57 template <class T>
59  m_root = 0;
60 }
61 
62 template <class T>
64  if (m_root) delete m_root;
65 }
66 
67 template <class T>
68 bool DiscretePDF<T>::empty() const {
69  return m_root == 0;
70 }
71 
72 template <class T>
73 void DiscretePDF<T>::insert(T item, weight_type weight) {
74  Node *p=0, *n=m_root;
75 
76  while (n) {
77  if (!n->leftIsBlack() && !n->rightIsBlack())
78  split(n);
79 
80  p = n;
81  if (n->key==item) {
82  assert(0 && "insert: argument(item) already in tree");
83  } else {
84  n = (item<n->key)?n->left:n->right;
85  }
86  }
87 
88  n = new Node(item, weight, p);
89 
90  if (!p) {
91  m_root = n;
92  } else {
93  if (item<p->key) {
94  p->left = n;
95  } else {
96  p->right = n;
97  }
98 
99  split(n);
100  }
101 
102  propogateSumsUp(n);
103 }
104 
105 template <class T>
106 void DiscretePDF<T>::remove(T item) {
107  Node **np = lookup(item, 0);
108  Node *child, *n = *np;
109 
110  if (!n) {
111  assert(0 && "remove: argument(item) not in tree");
112  } else {
113  if (n->left) {
114  Node **leftMaxp = &n->left;
115 
116  while ((*leftMaxp)->right)
117  leftMaxp = &(*leftMaxp)->right;
118 
119  n->key = (*leftMaxp)->key;
120  n->weight = (*leftMaxp)->weight;
121 
122  np = leftMaxp;
123  n = *np;
124  }
125 
126  // node now has at most one child
127 
128  child = n->left?n->left:n->right;
129  *np = child;
130 
131  if (child) {
132  child->parent = n->parent;
133 
134  if (n->isBlack()) {
135  lengthen(child);
136  }
137  }
138 
139  propogateSumsUp(n->parent);
140 
141  n->left = n->right = 0;
142  delete n;
143  }
144 }
145 
146 template <class T>
147 void DiscretePDF<T>::update(T item, weight_type weight) {
148  Node *n = *lookup(item, 0);
149 
150  if (!n) {
151  assert(0 && "update: argument(item) not in tree");
152  } else {
153  n->weight = weight;
154  propogateSumsUp(n);
155  }
156 }
157 
158 template <class T>
159 T DiscretePDF<T>::choose(double p) {
160  if (p<0.0 || p>=1.0) {
161  assert(0 && "choose: argument(p) outside valid range");
162  } else if (!m_root) {
163  assert(0 && "choose: choose() called on empty tree");
164  } else {
165  weight_type w = (weight_type) (m_root->sumWeights * p);
166  Node *n = m_root;
167 
168  while (1) {
169  if (n->left) {
170  if (w<n->left->sumWeights) {
171  n = n->left;
172  continue;
173  } else {
174  w -= n->left->sumWeights;
175  }
176  }
177  if (w<n->weight || !n->right) {
178  break; // !n->right condition shouldn't be necessary, just sanity check
179  }
180  w -= n->weight;
181  n = n->right;
182  }
183 
184  return n->key;
185  }
186 }
187 
188 template <class T>
189 bool DiscretePDF<T>::inTree(T item) {
190  Node *n = *lookup(item, 0);
191 
192  return !!n;
193 }
194 
195 template <class T>
196 typename DiscretePDF<T>::weight_type DiscretePDF<T>::getWeight(T item) {
197  Node *n = *lookup(item, 0);
198  assert(n);
199  return n->weight;
200 }
201 
202 //
203 
204 template <class T>
205 typename DiscretePDF<T>::Node **
206 DiscretePDF<T>::lookup(T item, Node **parent_out) {
207  Node *n, *p=0, **np=&m_root;
208 
209  while ((n = *np)) {
210  if (n->key==item) {
211  break;
212  } else {
213  p = n;
214  if (item<n->key) {
215  np = &n->left;
216  } else {
217  np = &n->right;
218  }
219  }
220  }
221 
222  if (parent_out)
223  *parent_out = p;
224  return np;
225 }
226 
227 template <class T>
228 void DiscretePDF<T>::split(Node *n) {
229  if (n->left) n->left->markBlack();
230  if (n->right) n->right->markBlack();
231 
232  if (n->parent) {
233  Node *p = n->parent;
234 
235  n->markRed();
236 
237  if (!p->isBlack()) {
238  p->parent->markRed();
239 
240  // not same direction
241  if (!((n==p->left && p==p->parent->left) ||
242  (n==p->right && p==p->parent->right))) {
243  rotate(n);
244  p = n;
245  }
246 
247  rotate(p);
248  p->markBlack();
249  }
250  }
251 }
252 
253 template <class T>
254 void DiscretePDF<T>::rotate(Node *n) {
255  Node *p=n->parent, *pp=p->parent;
256 
257  n->parent = pp;
258  p->parent = n;
259 
260  if (n==p->left) {
261  p->left = n->right;
262  n->right = p;
263  if (p->left) p->left->parent = p;
264  } else {
265  p->right = n->left;
266  n->left = p;
267  if (p->right) p->right->parent = p;
268  }
269 
270  n->setSum();
271  p->setSum();
272 
273  if (!pp) {
274  m_root = n;
275  } else {
276  if (p==pp->left) {
277  pp->left = n;
278  } else {
279  pp->right = n;
280  }
281  }
282 }
283 
284 template <class T>
285 void DiscretePDF<T>::lengthen(Node *n) {
286  if (!n->isBlack()) {
287  n->markBlack();
288  } else if (n->parent) {
289  Node *sibling = n->sibling();
290 
291  if (sibling && !sibling->isBlack()) {
292  n->parent->markRed();
293  sibling->markBlack();
294 
295  rotate(sibling); // node sibling is now old sibling child, must be black
296  sibling = n->sibling();
297  }
298 
299  // sibling is black
300 
301  if (!sibling) {
302  lengthen(n->parent);
303  } else if (sibling->leftIsBlack() && sibling->rightIsBlack()) {
304  if (n->parent->isBlack()) {
305  sibling->markRed();
306  lengthen(n->parent);
307  } else {
308  sibling->markRed();
309  n->parent->markBlack();
310  }
311  } else {
312  if (n==n->parent->left && sibling->rightIsBlack()) {
313  rotate(sibling->left); // sibling->left must be red
314  sibling->markRed();
315  sibling->parent->markBlack();
316  sibling = sibling->parent;
317  } else if (n==n->parent->right && sibling->leftIsBlack()) {
318  rotate(sibling->right); // sibling->right must be red
319  sibling->markRed();
320  sibling->parent->markBlack();
321  sibling = sibling->parent;
322  }
323 
324  // sibling is black, and sibling's far child is red
325 
326  rotate(sibling);
327  if (!n->parent->isBlack())
328  sibling->markRed();
329  sibling->left->markBlack();
330  sibling->right->markBlack();
331  }
332  }
333 }
334 
335 template <class T>
336 void DiscretePDF<T>::propogateSumsUp(Node *n) {
337  for (; n; n=n->parent)
338  n->setSum();
339 }
340 
341 }
342 
void insert(T item, weight_type weight)
void rotate(Node *node)
bool inTree(T item)
void remove(T item)
void propogateSumsUp(Node *n)
Node ** lookup(T item, Node **parent_out)
void split(Node *node)
void lengthen(Node *node)
T choose(double p)
void update(T item, weight_type newWeight)
weight_type getWeight(T item)
bool empty() const