1/* ----------------------------------------------------------------------------
2
3 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
4 * Atlanta, Georgia 30332-0415
5 * All Rights Reserved
6 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)
7
8 * See LICENSE for the license information
9
10 * -------------------------------------------------------------------------- */
11
12/**
13 * @file BTree.h
14 * @brief purely functional binary tree
15 * @author Chris Beall
16 * @author Frank Dellaert
17 * @date Feb 3, 2010
18 */
19
20#pragma once
21
22#include <stack>
23#include <sstream>
24#include <memory>
25#include <functional>
26
27namespace gtsam {
28
29 /**
30 * @brief Binary tree
31 * @ingroup base
32 */
33 template<class KEY, class VALUE>
34 class BTree {
35
36 public:
37
38 typedef std::pair<KEY, VALUE> value_type;
39
40 private:
41
42 /**
43 * Node in a tree
44 */
45 struct Node {
46
47 const size_t height_;
48 const value_type keyValue_;
49 const BTree left, right;
50
51 /** default constructor */
52 Node() {
53 }
54
55 /**
56 * Create leaf node with height 1
57 * @param keyValue (key,value) pair
58 */
59 Node(const value_type& keyValue) :
60 height_(1), keyValue_(keyValue) {
61 }
62
63 /**
64 * Create a node from two subtrees and a key value pair
65 */
66 Node(const BTree& l, const value_type& keyValue, const BTree& r) :
67 height_(l.height() >= r.height() ? l.height() + 1 : r.height() + 1),
68 keyValue_(keyValue), left(l), right(r) {
69 }
70
71 inline const KEY& key() const { return keyValue_.first;}
72 inline const VALUE& value() const { return keyValue_.second;}
73
74 }; // Node
75
76 // We store a shared pointer to the root of the functional tree
77 // composed of Node classes. If root_==nullptr, the tree is empty.
78 typedef std::shared_ptr<const Node> sharedNode;
79 sharedNode root_;
80
81 inline const value_type& keyValue() const { return root_->keyValue_;}
82 inline const KEY& key() const { return root_->key(); }
83 inline const VALUE& value() const { return root_->value(); }
84 inline const BTree& left() const { return root_->left; }
85 inline const BTree& right() const { return root_->right; }
86
87 /** create a new balanced tree out of two trees and a key-value pair */
88 static BTree balance(const BTree& l, const value_type& xd, const BTree& r) {
89 size_t hl = l.height(), hr = r.height();
90 if (hl > hr + 2) {
91 const BTree& ll = l.left(), lr = l.right();
92 if (ll.height() >= lr.height())
93 return BTree(ll, l.keyValue(), BTree(lr, xd, r));
94 else {
95 BTree _left(ll, l.keyValue(), lr.left());
96 BTree _right(lr.right(), xd, r);
97 return BTree(_left, lr.keyValue(), _right);
98 }
99 } else if (hr > hl + 2) {
100 const BTree& rl = r.left(), rr = r.right();
101 if (rr.height() >= rl.height())
102 return BTree(BTree(l, xd, rl), r.keyValue(), rr);
103 else {
104 BTree _left(l, xd, rl.left());
105 BTree _right(rl.right(), r.keyValue(), rr);
106 return BTree(_left, rl.keyValue(), _right);
107 }
108 } else
109 return BTree(l, xd, r);
110 }
111
112 public:
113
114 /** default constructor creates an empty tree */
115 BTree() {
116 }
117
118 /** copy constructor */
119 BTree(const BTree& other) :
120 root_(other.root_) {
121 }
122
123 /** create leaf from key-value pair */
124 BTree(const value_type& keyValue) :
125 root_(new Node(keyValue)) {
126 }
127
128 /** create from key-value pair and left, right subtrees */
129 BTree(const BTree& l, const value_type& keyValue, const BTree& r) :
130 root_(new Node(l, keyValue, r)) {
131 }
132
133 /** assignment operator */
134 BTree & operator= (const BTree & other) {
135 root_ = other.root_;
136 return *this;
137 }
138
139 /** Check whether tree is empty */
140 bool empty() const {
141 return !root_;
142 }
143
144 /** add a key-value pair */
145 BTree add(const value_type& xd) const {
146 if (empty()) return BTree(xd);
147 const KEY& x = xd.first;
148 if (x == key())
149 return BTree(left(), xd, right());
150 else if (x < key())
151 return balance(l: left().add(xd), xd: keyValue(), r: right());
152 else
153 return balance(l: left(), xd: keyValue(), r: right().add(xd));
154 }
155
156 /** add a key-value pair */
157 BTree add(const KEY& x, const VALUE& d) const {
158 return add(std::make_pair(x, d));
159 }
160
161 /** member predicate */
162 bool mem(const KEY& x) const {
163 if (!root_) return false;
164 if (x == key()) return true;
165 if (x < key())
166 return left().mem(x);
167 else
168 return right().mem(x);
169 }
170
171 /** Check whether trees are *exactly* the same (occupy same memory) */
172 inline bool same(const BTree& other) const {
173 return (other.root_ == root_);
174 }
175
176 /**
177 * Check whether trees are structurally the same,
178 * i.e., contain the same values in same tree-structure.
179 */
180 bool operator==(const BTree& other) const {
181 if (other.root_ == root_) return true; // if same, we're done
182 if (empty() && !other.empty()) return false;
183 if (!empty() && other.empty()) return false;
184 // both non-empty, recurse: check this key-value pair and subtrees...
185 return (keyValue() == other.keyValue()) && (left() == other.left())
186 && (right() == other.right());
187 }
188
189 inline bool operator!=(const BTree& other) const {
190 return !operator==(other);
191 }
192
193 /** minimum key binding */
194 const value_type& min() const {
195 if (!root_) throw std::invalid_argument("BTree::min: empty tree");
196 if (left().empty()) return keyValue();
197 return left().min();
198 }
199
200 /** remove minimum key binding */
201 BTree remove_min() const {
202 if (!root_) throw std::invalid_argument("BTree::remove_min: empty tree");
203 if (left().empty()) return right();
204 return balance(l: left().remove_min(), xd: keyValue(), r: right());
205 }
206
207 /** merge two trees */
208 static BTree merge(const BTree& t1, const BTree& t2) {
209 if (t1.empty()) return t2;
210 if (t2.empty()) return t1;
211 const value_type& xd = t2.min();
212 return balance(l: t1, xd, r: t2.remove_min());
213 }
214
215 /** remove a key-value pair */
216 BTree remove(const KEY& x) const {
217 if (!root_) return BTree();
218 if (x == key())
219 return merge(t1: left(), t2: right());
220 else if (x < key())
221 return balance(l: left().remove(x), xd: keyValue(), r: right());
222 else
223 return balance(l: left(), xd: keyValue(), r: right().remove(x));
224 }
225
226 /** Return height of the tree, 0 if empty */
227 size_t height() const {
228 return (root_ != nullptr) ? root_->height_ : 0;
229 }
230
231 /** return size of the tree */
232 size_t size() const {
233 if (!root_) return 0;
234 return left().size() + 1 + right().size();
235 }
236
237 /**
238 * find a value given a key, throws exception when not found
239 * Optimized non-recursive version as [find] is crucial for speed
240 */
241 const VALUE& find(const KEY& k) const {
242 const Node* node = root_.get();
243 while (node) {
244 const KEY& key = node->key();
245 if (k < key) node = node->left.root_.get();
246 else if (key < k) node = node->right.root_.get();
247 else return node->value();
248 }
249
250 throw std::invalid_argument("BTree::find: key not found");
251 }
252
253 /** print in-order */
254 void print(const std::string& s = "") const {
255 if (empty()) return;
256 KEY k = key();
257 std::stringstream ss;
258 ss << height();
259 k.print(s + ss.str() + " ");
260 left().print(s + "L ");
261 right().print(s + "R ");
262 }
263
264 /** iterate over tree */
265 void iter(std::function<void(const KEY&, const VALUE&)> f) const {
266 if (!root_) return;
267 left().iter(f);
268 f(key(), value());
269 right().iter(f);
270 }
271
272 /** map key-values in tree over function f that computes a new value */
273 template<class TO>
274 BTree<KEY, TO> map(std::function<TO(const KEY&, const VALUE&)> f) const {
275 if (empty()) return BTree<KEY, TO> ();
276 std::pair<KEY, TO> xd(key(), f(key(), value()));
277 return BTree<KEY, TO> (left().map(f), xd, right().map(f));
278 }
279
280 /**
281 * t.fold(f,a) computes [(f kN dN ... (f k1 d1 a)...)],
282 * where [k1 ... kN] are the keys of all bindings in [m],
283 * and [d1 ... dN] are the associated data.
284 * The associated values are passed to [f] in reverse sort order
285 */
286 template<class ACC>
287 ACC fold(std::function<ACC(const KEY&, const VALUE&, const ACC&)> f,
288 const ACC& a) const {
289 if (!root_) return a;
290 ACC ar = right().fold(f, a); // fold over right subtree
291 ACC am = f(key(), value(), ar); // apply f with current value
292 return left().fold(f, am); // fold over left subtree
293 }
294
295 /**
296 * @brief Const iterator
297 * Not trivial: iterator keeps a stack to indicate current path from root_
298 */
299 class const_iterator {
300
301 private:
302
303 typedef const_iterator Self;
304 typedef std::pair<sharedNode, bool> flagged;
305
306 /** path to the iterator, annotated with flag */
307 std::stack<flagged> path_;
308
309 const sharedNode& current() const {
310 return path_.top().first;
311 }
312
313 bool done() const {
314 return path_.top().second;
315 }
316
317 // The idea is we already iterated through the left-subtree and current key-value.
318 // We now try pushing left subtree of right onto the stack. If there is no right
319 // sub-tree, we pop this node of the stack and the parent becomes the iterator.
320 // We avoid going down a right-subtree that was already visited by checking the flag.
321 void increment() {
322 if (path_.empty()) return;
323 sharedNode t = current()->right.root_;
324 if (!t || done()) {
325 // no right subtree, iterator becomes first parent with a non-visited right subtree
326 path_.pop();
327 while (!path_.empty() && done())
328 path_.pop();
329 } else {
330 path_.top().second = true; // flag we visited right
331 // push right root and its left-most path onto the stack
332 while (t) {
333 path_.push(std::make_pair(t, false));
334 t = t->left.root_;
335 }
336 }
337 }
338
339 public:
340
341 // traits for playing nice with STL
342 typedef ptrdiff_t difference_type;
343 typedef std::forward_iterator_tag iterator_category;
344 typedef std::pair<KEY, VALUE> value_type;
345 typedef const value_type* pointer;
346 typedef const value_type& reference;
347
348 /** initialize end */
349 const_iterator() {
350 }
351
352 /** initialize from root */
353 const_iterator(const sharedNode& root) {
354 sharedNode t = root;
355 while (t) {
356 path_.push(std::make_pair(t, false));
357 t = t->left.root_;
358 }
359 }
360
361 /** equality */
362 bool operator==(const Self& __x) const {
363 return path_ == __x.path_;
364 }
365
366 /** inequality */
367 bool operator!=(const Self& __x) const {
368 return path_ != __x.path_;
369 }
370
371 /** dereference */
372 reference operator*() const {
373 if (path_.empty()) throw std::invalid_argument(
374 "operator*: tried to dereference end");
375 return current()->keyValue_;
376 }
377
378 /** dereference */
379 pointer operator->() const {
380 if (path_.empty()) throw std::invalid_argument(
381 "operator->: tried to dereference end");
382 return &(current()->keyValue_);
383 }
384
385 /** pre-increment */
386 Self& operator++() {
387 increment();
388 return *this;
389 }
390
391 /** post-increment */
392 Self operator++(int) {
393 Self __tmp = *this;
394 increment();
395 return __tmp;
396 }
397
398 }; // const_iterator
399
400 // to make BTree work with range-based for
401 // We do *not* want a non-const iterator
402 typedef const_iterator iterator;
403
404 /** return iterator */
405 const_iterator begin() const {
406 return const_iterator(root_);
407 }
408
409 /** return iterator */
410 const_iterator end() const {
411 return const_iterator();
412 }
413
414 }; // BTree
415
416} // namespace gtsam
417
418