1 /* 2 Copyright (C) 2005-2007 Feeling Software Inc. 3 Portions of the code are: 4 Copyright (C) 2005-2007 Sony Computer Entertainment America 5 6 MIT License: http://www.opensource.org/licenses/mit-license.php 7 */ 8 9 /** 10 @file FMTree.h 11 The file contains the tree class, which improves on the standard C++ tree class. 12 */ 13 14 #ifndef _FM_TREE_H_ 15 #define _FM_TREE_H_ 16 17 #ifndef _FM_ALLOCATOR_H_ 18 #include "FMath/FMAllocator.h" 19 #endif // _FM_ALLOCATOR_H_ 20 21 namespace fm 22 { 23 /** 24 A simple pair template. 25 @ingroup FMath 26 */ 27 template <class _Kty, class _Ty> 28 class pair 29 { 30 public: 31 _Kty first; /**< The first element of the pair. */ 32 _Ty second; /**< The second element of the pair. */ 33 34 /** Constructor. */ pair()35 pair() : first(), second() {} 36 37 /** Constructor. 38 @param f A first element to copy. 39 @param s A second element to copy. */ pair(const _Kty & f,const _Ty & s)40 pair(const _Kty& f, const _Ty& s) : first(f), second(s) {} 41 42 /** Copy constructor. 43 @param p A pair to clone. */ pair(const pair & p)44 pair(const pair& p) : first(p.first), second(p.second) {} 45 46 /** Returns whether the two pairs are equal. 47 @param p A second pair of the same type. 48 @return Whether the two pairs are equal. */ 49 inline bool operator==(const pair& p) const { return p.first == first && p.second == second; } 50 51 /** Returns whether the two pairs are not equal. 52 @param p A second pair of the same type. 53 @return Whether the two pairs are not equal. */ 54 inline bool operator!=(const pair& p) const { return p.first != first || p.second != second; } 55 }; 56 57 /** 58 An auto-balancing tree. 59 Intentionally has an interface similar to the standard C++ tree class. 60 It's implement should be very similar yet more lightweight. 61 62 @ingroup FMath 63 */ 64 template <class KEY, class DATA> 65 class tree 66 { 67 private: 68 typedef fm::pair<KEY, DATA> pair; 69 70 class node 71 { 72 public: 73 node* left; 74 node* right; 75 node* parent; 76 77 int32 weight; 78 79 pair data; 80 81 public: node()82 node() : left(NULL), right(NULL), parent(NULL), weight(0) {} 83 rotateLeft()84 void rotateLeft() 85 { 86 node** parentLink = (parent->left == this) ? &parent->left : &parent->right; 87 88 // detach right 89 node* oldRight = right; 90 91 // detach right's left and attach to the parent's right. 92 node* right_left = right->left; 93 right = right_left; 94 if (right_left != NULL) right_left->parent = this; 95 96 // attach the right to the double parent. 97 oldRight->left = this; 98 99 // attach the parent on the right's left. 100 oldRight->parent = parent; 101 parent = oldRight; 102 (*parentLink) = oldRight; 103 104 // adjust the weights 105 weight = weight - 1 - (oldRight->weight > 0 ? oldRight->weight : 0); 106 oldRight->weight = oldRight->weight - (0 > -weight ? 0 : -weight) - 1; 107 } 108 rotateRight()109 void rotateRight() 110 { 111 node** parentLink = (parent->left == this) ? &parent->left : &parent->right; 112 113 // detach left 114 node* oldLeft = left; 115 116 // detach left's right and attach to the parent's left. 117 node* left_right = left->right; 118 left = left_right; 119 if (left_right != NULL) left_right->parent = this; 120 121 // attach the parent on the left's right. 122 oldLeft->right = this; 123 124 // attach the left to the double parent. 125 oldLeft->parent = parent; 126 parent = oldLeft; 127 (*parentLink) = oldLeft; 128 129 // adjust the weights 130 weight = weight + 1 + (0 > oldLeft->weight ? -oldLeft->weight : 0); 131 oldLeft->weight = (oldLeft->weight + 1) + (0 > weight ? 0 : weight); 132 } 133 134 #ifdef TREE_DEBUG depth()135 intptr_t depth() const 136 { 137 intptr_t leftDepth = left != NULL ? left->depth() : 0; 138 intptr_t rightDepth = right != NULL ? right->depth() : 0; 139 return max(leftDepth, rightDepth) + 1; 140 } 141 is_correct()142 void is_correct() 143 { 144 if (left != NULL) left->is_correct(); 145 if (right != NULL) right->is_correct(); 146 intptr_t leftDepth = left != NULL ? left->depth() : 0; 147 intptr_t rightDepth = right != NULL ? right->depth() : 0; 148 FUAssert(rightDepth - leftDepth == weight,); 149 FUAssert(abs(weight) < 2,); 150 } 151 #endif // TREE_DEBUG 152 }; 153 154 public: 155 class const_iterator; 156 157 /** 158 A tree element iterator. 159 Similar to the basic STL iterator. 160 */ 161 class iterator 162 { 163 private: 164 friend class tree; 165 friend class const_iterator; 166 node* currentNode; 167 168 public: 169 /** Empty constructor. */ iterator()170 iterator() {} 171 /** Constructor. @param n The tree node at which to start the iteration. */ iterator(node * n)172 iterator(node* n) : currentNode(n) {} 173 /** Copy operator. @param copy The tree iterator to copy. */ 174 iterator& operator=(const iterator& copy) { currentNode = copy.currentNode; return *this; } 175 176 /** Retrieves whether this iterator points to the same node as the given iterator. 177 @param other A second iterator. 178 @return Whether the two iterators are pointing to the same node. */ 179 inline bool operator==(const iterator& other) const { return other.currentNode == currentNode; } 180 inline bool operator==(const const_iterator& other) const; /**< See above. */ 181 182 /** Retrieves whether this iterator points to a different node that a given iterator. 183 @param other A second iterator. 184 @return Whether the two iterators are pointing at different nodes. */ 185 inline bool operator!=(const iterator& other) const { return other.currentNode != currentNode; } 186 inline bool operator!=(const const_iterator& other) const;/**< See above. */ 187 188 /** Advances the iterator to the next ordered tree node. 189 @return This iterator. */ 190 iterator& operator++() 191 { 192 // Go one right or one up. 193 if (currentNode->right == NULL) 194 { 195 node* oldNode; 196 do 197 { 198 oldNode = currentNode; 199 200 // Go one up. 201 // We control the root node, which is the only with parent == NULL. 202 // if you crash here, you don't check for iterator == end() correctly. 203 currentNode = currentNode->parent; 204 } 205 while (currentNode->right == oldNode && currentNode->parent != NULL); 206 } 207 else 208 { 209 // Go one right. 210 currentNode = currentNode->right; 211 212 // Go all the way left. 213 while (currentNode->left != NULL) currentNode = currentNode->left; 214 } 215 return (*this); 216 } 217 218 /** Backtrack the iterator to the next ordered tree node. 219 @return This iterator. */ 220 iterator& operator--() 221 { 222 // Go one left or one up. 223 if (currentNode->left == NULL) 224 { 225 node* oldNode; 226 do 227 { 228 oldNode = currentNode; 229 230 // Go one up. 231 // We control the root node, which is the only with parent == NULL. 232 // if you crash here, you don't check for iterator == begin() correctly. 233 currentNode = currentNode->parent; 234 } 235 while (currentNode->left == oldNode && currentNode->parent != NULL); 236 } 237 else 238 { 239 // Go one left. 240 currentNode = currentNode->left; 241 242 // Go all the way right. 243 while (currentNode->right != NULL) currentNode = currentNode->right; 244 } 245 return (*this); 246 } 247 248 /** Retrieves the current tree node. 249 @return The current tree node. */ 250 inline pair& operator*() { return currentNode->data; } 251 inline pair* operator->() { return ¤tNode->data; } /**< See above. */ 252 }; 253 254 /** 255 A tree constant-element iterator. 256 Similar to the basic STL const_iterator. 257 */ 258 class const_iterator 259 { 260 private: 261 friend class tree; 262 friend class iterator; 263 const node* currentNode; 264 265 public: 266 /** Empty constructor. */ const_iterator()267 const_iterator() {} 268 /** Copy constructor. 269 @param copy The iterator to clone. */ const_iterator(const iterator & copy)270 const_iterator(const iterator& copy) : currentNode(copy.currentNode) {} 271 /** Constructor. @param n The tree node at which to start the iteration. */ const_iterator(const node * n)272 const_iterator(const node* n) : currentNode(n) {} 273 /** Copy operator. @param copy The tree iterator to copy. */ 274 const_iterator& operator=(const iterator& copy) { currentNode = copy.currentNode; return *this; } 275 const_iterator& operator=(const const_iterator& copy) { currentNode = copy.currentNode; return *this; } /**< See above. */ 276 277 /** Retrieves whether this iterator points to the same node as the given iterator. 278 @param other A second iterator. 279 @return Whether the two iterators are pointing to the same node. */ 280 inline bool operator==(const iterator& other) const { return other.currentNode == currentNode; } 281 inline bool operator==(const const_iterator& other) const { return other.currentNode == currentNode; } /**< See above. */ 282 283 /** Retrieves whether this iterator points to a different node that a given iterator. 284 @param other A second iterator. 285 @return Whether the two iterators are pointing at different nodes. */ 286 inline bool operator!=(const iterator& other) const { return other.currentNode != currentNode; } 287 inline bool operator!=(const const_iterator& other) const { return other.currentNode != currentNode; } /**< See above. */ 288 289 /** Advances the iterator to the next ordered tree node. 290 @return This iterator. */ 291 const_iterator& operator++() 292 { 293 // Go one right or one up. 294 if (currentNode->right == NULL) 295 { 296 const node* oldNode; 297 do 298 { 299 oldNode = currentNode; 300 301 // Go one up. 302 // We control the root node, which is the only with parent == NULL. 303 // if you crash here, you don't check for iterator == end() correctly. 304 currentNode = currentNode->parent; 305 } 306 while (currentNode->right == oldNode && currentNode->parent != NULL); 307 } 308 else 309 { 310 // Go one right. 311 currentNode = currentNode->right; 312 313 // Go all the way left. 314 while (currentNode->left != NULL) currentNode = currentNode->left; 315 } 316 return (*this); 317 } 318 319 /** Backtrack the iterator to the next ordered tree node. 320 @return This iterator. */ 321 const_iterator& operator--() 322 { 323 // Go one left or one up. 324 if (currentNode->left == NULL) 325 { 326 const node* oldNode; 327 do 328 { 329 oldNode = currentNode; 330 331 // Go one up. 332 // We control the root node, which is the only with parent == NULL. 333 // if you crash here, you don't check for iterator == end() correctly. 334 currentNode = currentNode->parent; 335 } 336 while (currentNode->left == oldNode && currentNode->parent != NULL); 337 } 338 else 339 { 340 // Go one left. 341 currentNode = currentNode->left; 342 343 // Go all the way right. 344 while (currentNode->right != NULL) currentNode = currentNode->right; 345 } 346 return (*this); 347 } 348 349 /** Retrieves the current tree node. 350 @return The current tree node. */ 351 inline const pair& operator*() { return currentNode->data; } 352 inline const pair* operator->() { return ¤tNode->data; } /**< See above. */ 353 }; 354 355 private: 356 node* root; 357 size_t sized; 358 359 public: 360 /** Constructor. */ tree()361 tree() : root(NULL), sized(0) 362 { 363 root = (node*) fm::Allocate(sizeof(node)); 364 fm::Construct(root); 365 } 366 367 /** Destructor. */ ~tree()368 ~tree() 369 { 370 clear(); 371 root->data.first.~KEY(); 372 root->data.second.~DATA(); 373 fm::Release(root); 374 root = NULL; 375 } 376 377 /** Retrieves the first ordered element within the tree. 378 @return An iterator that points to the first tree element. */ begin()379 inline iterator begin() { iterator it(root); return (root->right == NULL) ? it : ++it; } begin()380 inline const_iterator begin() const { const_iterator it(root); return (root->right == NULL) ? it : ++it; } /**< See above. */ 381 382 /** Retrieves an iterator that points just passed the last 383 ordered element within the tree. 384 @return An iterator just passed the last element in the tree. */ end()385 inline iterator end() { return iterator(root); } end()386 inline const_iterator end() const { return const_iterator(root); } /**< See above. */ 387 388 /** Retrieves the last ordered element within the tree. 389 @return An iterator that points to the last tree element. */ last()390 inline iterator last() { node* n = root; while (n->right != NULL) n = n->right; return iterator(n); } last()391 inline const_iterator last() const { const node* n = root; while (n->right != NULL) n = n->right; return const_iterator(n); } /**< See above. */ 392 393 /** Retrieves an existing data element using its key. 394 @param key The key. 395 @return An iterator that points to the existing data element. 396 This iterator may be past the end of the tree, in the case 397 where the key does not exists within the tree: do check the 398 iterator against end(). */ find(const KEY & key)399 iterator find(const KEY& key) 400 { 401 node* out = root->right; 402 while (out != NULL) 403 { 404 if (key < out->data.first) out = out->left; 405 else if (key == out->data.first) return iterator(out); 406 else out = out->right; 407 } 408 return end(); 409 } find(const KEY & key)410 inline const_iterator find(const KEY& key) const { return const_iterator(const_cast<tree<KEY, DATA>* >(this)->find(key)); } /**< See above. */ 411 412 /** Inserts a new data element with its key. 413 If the key already exists within the tree, the 414 old data element will be overwritten using the new data element. 415 @param key The new key. 416 @param data The new data element. 417 @return An iterator that points to the tree element. */ insert(const KEY & key,const DATA & data)418 iterator insert(const KEY& key, const DATA& data) 419 { 420 // First step: look for an already existing entry. 421 node** insertAt = &root->right,* parent = root; 422 while (*insertAt != NULL) 423 { 424 parent = *insertAt; 425 if (key < parent->data.first) insertAt = &parent->left; 426 else if (key == parent->data.first) 427 { 428 parent->data.second = data; 429 return iterator(parent); 430 } 431 else insertAt = &parent->right; 432 } 433 434 // Insert the new node. 435 node* n = ((*insertAt) = (node*) fm::Allocate(sizeof(node))); 436 fm::Construct(*insertAt); // could be get rid of this one? it would work for all pointer and primitive types.. 437 n->parent = parent; 438 n->data.first = key; 439 n->data.second = data; 440 ++sized; 441 442 // Balance the tree. 443 parent->weight += (*insertAt == parent->right) ? 1 : -1; 444 node* it = parent; 445 while (it != root) 446 { 447 // Check whether we need to balance this level. 448 if (it->weight > 1) 449 { 450 if (it->right->weight < 0) it->right->rotateRight(); 451 it->rotateLeft(); 452 it = it->parent; 453 break; 454 } 455 else if (it->weight < -1) 456 { 457 if (it->left->weight > 0) it->left->rotateLeft(); 458 it->rotateRight(); 459 it = it->parent; 460 break; 461 } 462 else if (it->weight == 0) break; // no height change. 463 464 // go up one level. 465 it->parent->weight += (it == it->parent->right) ? 1 : -1; 466 it = it->parent; 467 } 468 469 #ifdef TREE_DEBUG 470 root->right->is_correct(); 471 #endif // TREE_DEBUG 472 473 return iterator(n); 474 } 475 476 /** Retrieves a data element using its key. 477 @param k The key. 478 @return The data element for this key. In the non-constant 479 version of this function, a new element is created for 480 the key if it does not already belong to the tree. */ 481 inline DATA& operator[](const KEY& k) { iterator it = find(k); if (it != end()) return it->second; else { DATA d; return insert(k, d)->second; } } 482 inline const DATA& operator[](const KEY& k) const { return find(k)->second; } /**< See above. */ 483 484 /** Removes a data element from the tree. 485 @param key The key of the data element to erase. */ erase(const KEY & key)486 inline void erase(const KEY& key) { iterator it = find(key); erase(it); } 487 488 /** Removes a data element from the tree. 489 @param it An iterator that points to the tree element to erase. */ erase(const iterator & it)490 inline void erase(const iterator& it) 491 { 492 node* n = it.currentNode; 493 if (n != root) 494 { 495 node* release; 496 if (n->left == NULL && n->right == NULL) release = n; 497 else 498 { 499 // choose whether to reduce on the left or right. 500 if (n->weight <= 0 && n->left != NULL) 501 { 502 // take out the left's rightmost node. 503 release = n->left; 504 while (release->right != NULL) release = release->right; 505 n->data = release->data; 506 507 // push up any left node on the rightmost node. 508 if (release->left != NULL) 509 { 510 release->data = release->left->data; 511 release = release->left; 512 } 513 } 514 else 515 { 516 // take out the right's leftmost node. 517 release = n->right; 518 while (release->left != NULL) release = release->left; 519 n->data = release->data; 520 521 // push up any right node on the leftmost node. 522 if (release->right != NULL) 523 { 524 release->data = release->right->data; 525 release = release->right; 526 } 527 } 528 } 529 530 // Release the selected node and re-adjust its parent's weight. 531 node* rebalance = release->parent; 532 if (rebalance->left == release) { rebalance->left = NULL; ++rebalance->weight; } 533 else { rebalance->right = NULL; --rebalance->weight; } 534 release->data.first.~KEY(); 535 release->data.second.~DATA(); 536 fm::Release(release); 537 --sized; 538 539 // Rebalance the tree. 540 node* it = rebalance; 541 while (it != root) 542 { 543 // Check whether we need to balance this level. 544 if (it->weight > 1) 545 { 546 if (it->right->weight < 0) it->right->rotateRight(); 547 it->rotateLeft(); 548 it = it->parent; 549 } 550 else if (it->weight < -1) 551 { 552 if (it->left->weight > 0) it->left->rotateLeft(); 553 it->rotateRight(); 554 it = it->parent; 555 } 556 if (it->weight != 0) break; 557 558 // go up one level. 559 it->parent->weight -= (it == it->parent->right) ? 1 : -1; 560 it = it->parent; 561 } 562 } 563 564 #ifdef TREE_DEBUG 565 if (root->right != NULL) root->right->is_correct(); 566 #endif // TREE_DEBUG 567 } 568 569 /** Retrieves whether the tree contains any data nodes. 570 @return Whether the tree contains any data nodes. */ empty()571 inline bool empty() const { return sized == 0; } 572 573 /** Retrieves the number of data nodes contained in the tree. 574 @return The number of data nodes contained in the tree. */ size()575 inline size_t size() const { return sized; } 576 577 /** Removes all the data nodes from the tree. 578 This effectively prunes at the tree root. */ clear()579 void clear() 580 { 581 // Need to delete all the nodes. 582 if (root->right != NULL) 583 { 584 node* n = root->right; 585 while (n != root) 586 { 587 if (n->left != NULL) n = n->left; 588 else if (n->right != NULL) n = n->right; 589 else 590 { 591 // destroy this node. 592 node* release = n; 593 n = n->parent; 594 if (n->left == release) n->left = NULL; 595 else if (n->right == release) n->right = NULL; 596 release->data.first.~KEY(); 597 release->data.second.~DATA(); 598 fm::Release(release); 599 --sized; 600 } 601 } 602 root->right = NULL; 603 } 604 } 605 606 /** Copy constructor. Clones another tree into this one. 607 This is a bad function: it costs too much performance, avoid at all costs! 608 @param copy The tree to clone. 609 @return This tree, cloned. */ 610 inline tree<KEY,DATA>& operator= (const tree<KEY,DATA>& copy) 611 { 612 clear(); 613 614 // Function based on the iterator.. 615 // Go one right or one up. 616 node* currentNode = copy.root; 617 node* cloneNode = root; 618 if (currentNode->right != NULL) 619 { 620 do 621 { 622 if (currentNode->right == NULL) 623 { 624 const node* oldNode; 625 do 626 { 627 oldNode = currentNode; 628 629 // Go one up. 630 // We control the root node, which is the only with parent == NULL. 631 // if you crash here, you don't check for iterator == end() correctly. 632 currentNode = currentNode->parent; 633 cloneNode = cloneNode->parent; 634 } 635 while (currentNode->right == oldNode && currentNode->parent != NULL); 636 } 637 else 638 { 639 // Create and go one right. 640 currentNode = currentNode->right; 641 642 cloneNode->right = (node*) fm::Allocate(sizeof(node)); 643 fm::Construct(cloneNode->right); 644 cloneNode->right->parent = cloneNode; 645 cloneNode->right->data = currentNode->data; 646 cloneNode->right->weight = currentNode->weight; 647 ++sized; 648 649 cloneNode = cloneNode->right; 650 651 // Create and go one all the way left. 652 while (currentNode->left != NULL) 653 { 654 currentNode = currentNode->left; 655 656 cloneNode->left = (node*) fm::Allocate(sizeof(node)); 657 fm::Construct(cloneNode->left); 658 cloneNode->left->parent = cloneNode; 659 cloneNode->left->data = currentNode->data; 660 cloneNode->left->weight = currentNode->weight; 661 ++sized; 662 663 cloneNode = cloneNode->left; 664 } 665 } 666 } 667 while (currentNode != copy.root); 668 } 669 670 return (*this); 671 } 672 }; 673 674 template <class KEY, class DATA> 675 bool tree<KEY,DATA>::iterator::operator==(const const_iterator& other) const { return other.currentNode == currentNode; } 676 template <class KEY, class DATA> 677 bool tree<KEY,DATA>::iterator::operator!=(const const_iterator& other) const { return other.currentNode != currentNode; } 678 679 /** A STL set. */ 680 template <class _Kty> 681 class set : public fm::tree<_Kty, _Kty> {}; 682 683 /** A STL map. */ 684 template <class _Kty, class _Ty> 685 class map : public fm::tree<_Kty, _Ty> {}; 686 }; 687 688 #endif // _FM_TREE_H_ 689 690 691