wimlib.net Git - wimlib/blob - src/avl_tree.c

   1 /*
   2  * avl_tree.c
   3  *
   4  * Intrusive, nonrecursive AVL tree data structure (self-balancing binary search
   5  * tree), implementation file.
   6  *
   7  * Author:  Eric Biggers
   8  * Year:    2014
   9  *
  10  * This file is placed into the public domain.  You can do whatever you want
  11  * with it.
  12  */
  13
  14 #include <wimlib/avl_tree.h>
  15
  16 /* Starts an in-order traversal of the tree: returns the least-valued node, or
  17  * NULL if the tree is empty.  */
  18 struct avl_tree_node *
  19 avl_tree_first_in_order(const struct avl_tree_node *root)
  20 {
  21         const struct avl_tree_node *first = root;
  22
  23         if (first)
  24                 while (first->left)
  25                         first = first->left;
  26         return (struct avl_tree_node *)first;
  27 }
  28
  29 /* Continues an in-order traversal of the tree: returns the next-greatest-valued
  30  * node, or NULL if there is none.  */
  31 struct avl_tree_node *
  32 avl_tree_next_in_order(const struct avl_tree_node *prev)
  33 {
  34         const struct avl_tree_node *next;
  35
  36         if (prev->right)
  37                 for (next = prev->right;
  38                      next->left;
  39                      next = next->left)
  40                         ;
  41         else
  42                 for (next = avl_get_parent(prev);
  43                      next && prev == next->right;
  44                      prev = next, next = avl_get_parent(next))
  45                         ;
  46         return (struct avl_tree_node *)next;
  47 }
  48
  49 /* Starts a postorder traversal of the tree.  */
  50 struct avl_tree_node *
  51 avl_tree_first_in_postorder(const struct avl_tree_node *root)
  52 {
  53         const struct avl_tree_node *first = root;
  54
  55         if (first)
  56                 while (first->left || first->right)
  57                         first = first->left ? first->left : first->right;
  58
  59         return (struct avl_tree_node *)first;
  60 }
  61
  62 /* Continues a postorder traversal of the tree.  @prev will not be deferenced as
  63  * it's allowed that its memory has been freed; @prev_parent must be its saved
  64  * parent node.  Returns NULL if there are no more nodes (i.e. @prev was the
  65  * root of the tree).  */
  66 struct avl_tree_node *
  67 avl_tree_next_in_postorder(const struct avl_tree_node *prev,
  68                            const struct avl_tree_node *prev_parent)
  69 {
  70         const struct avl_tree_node *next = prev_parent;
  71
  72         if (next && prev == next->left && next->right)
  73                 for (next = next->right;
  74                      next->left || next->right;
  75                      next = next->left ? next->left : next->right)
  76                         ;
  77         return (struct avl_tree_node *)next;
  78 }
  79
  80 /* Get the left child (sign < 0) or the right child (sign > 0)
  81  * Note: for all calls of this, 'sign' is constant at compilation time and the
  82  * compiler can remove the conditional.  */
  83 static AVL_INLINE struct avl_tree_node *
  84 avl_get_child(const struct avl_tree_node *parent, int sign)
  85 {
  86         if (sign < 0)
  87                 return parent->left;
  88         else
  89                 return parent->right;
  90 }
  91
  92 /* Set the left child (sign < 0) or the right child (sign > 0)
  93  * Note: for all calls of this, 'sign' is constant at compilation time and the
  94  * compiler can remove the conditional.  */
  95 static AVL_INLINE void
  96 avl_set_child(struct avl_tree_node *parent, int sign,
  97               struct avl_tree_node *child)
  98 {
  99         if (sign < 0)
 100                 parent->left = child;
 101         else
 102                 parent->right = child;
 103 }
 104
 105 /* Sets the parent of the AVL tree @node to @parent.  */
 106 static AVL_INLINE void
 107 avl_set_parent(struct avl_tree_node *node, struct avl_tree_node *parent)
 108 {
 109         node->parent_balance =
 110                 (node->parent_balance & 3) | (uintptr_t)parent;
 111 }
 112
 113 /* Sets the parent and balance factor of the AVL tree @node.  */
 114 static AVL_INLINE void
 115 avl_set_parent_balance(struct avl_tree_node *node, struct avl_tree_node *parent,
 116                        int balance_factor)
 117 {
 118         node->parent_balance = (uintptr_t)parent | (balance_factor + 1);
 119 }
 120
 121 /* Returns the balance factor of the specified AVL tree node --- that is, the
 122  * height of its right subtree minus the height of its left subtree.  */
 123 static AVL_INLINE int
 124 avl_get_balance_factor(const struct avl_tree_node *node)
 125 {
 126         return (int)(node->parent_balance & 3) - 1;
 127 }
 128
 129 /* Sets the balance factor of the specified AVL tree node.  The caller MUST
 130  * ensure that this number is valid and maintains the AVL tree invariants.  */
 131 static AVL_INLINE void
 132 avl_set_balance_factor(struct avl_tree_node *node, int balance_factor)
 133 {
 134         node->parent_balance =
 135                 (node->parent_balance & ~3) | (balance_factor + 1);
 136 }
 137
 138 /* Increments the balance factor of the specified AVL tree @node by the
 139  * specified @amount.  The caller MUST ensure that this number is valid and
 140  * maintains the AVL tree invariants.  */
 141 static AVL_INLINE void
 142 avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
 143 {
 144         node->parent_balance += amount;
 145 }
 146
 147 /*
 148  * Template for performing a rotation ---
 149  *
 150  * Clockwise/Right (sign > 0):
 151  *
 152  *           X             X
 153  *           |             |
 154  *           A             B
 155  *          / \           / \
 156  *         B   C   =>    D   A
 157  *        / \               / \
 158  *       D   E             E   C
 159  *
 160  * Counterclockwise/Left (sign < 0)- same, except left and right are reversed:
 161  *
 162  *           X             X
 163  *           |             |
 164  *           A             B
 165  *          / \           / \
 166  *         C   B   =>    A   D
 167  *            / \       / \
 168  *           E   D     C   E
 169  *
 170  * This updates pointers but not balance factors!
 171  */
 172 static AVL_INLINE void
 173 avl_rotate(struct avl_tree_node ** const root_ptr,
 174            struct avl_tree_node * const A, const int sign)
 175 {
 176         struct avl_tree_node * const B = avl_get_child(A, -sign);
 177         struct avl_tree_node * const C = avl_get_child(A, +sign);
 178         struct avl_tree_node * const E = avl_get_child(B, +sign);
 179         struct avl_tree_node * const X = avl_get_parent(A);
 180
 181         avl_set_child(A, -sign, E);
 182         avl_set_child(A, +sign, C);
 183         avl_set_parent(A, B);
 184
 185         avl_set_child(B, +sign, A);
 186         avl_set_parent(B, X);
 187
 188         if (E)
 189                 avl_set_parent(E, A);
 190
 191         if (X) {
 192                 if (avl_get_child(X, +sign) == A)
 193                         avl_set_child(X, +sign, B);
 194                 else
 195                         avl_set_child(X, -sign, B);
 196         } else {
 197                 *root_ptr = B;
 198         }
 199 }
 200
 201 /*
 202  * Template for performing a double rotation ---
 203  *
 204  * Rotate counterclockwise (left) rooted at B, then clockwise (right) rooted at
 205  * A (sign > 0):
 206  *
 207  *           A             A           E
 208  *          / \           / \        /   \
 209  *         B   C   =>    E   C  =>  B     A
 210  *        / \           / \        / \   / \
 211  *       D   E         B   G      D   F G   C
 212  *          / \       / \
 213  *         F   G     D   F
 214  *
 215  * Rotate clockwise (right) rooted at B, then counterclockwise (left) rooted at
 216  * A (sign < 0):
 217  *
 218  *         A           A               E
 219  *        / \         / \            /   \
 220  *       C   B   =>  C   E    =>    A     B
 221  *          / \         / \        / \   / \
 222  *         E   D       G   B      C   G F   D
 223  *        / \             / \
 224  *       G   F           F   D
 225  *
 226  * Returns a pointer to E, the new root, and updates balance factors.  Except
 227  * for that, this is equivalent to:
 228  *      avl_rotate(root_ptr, B, -sign);
 229  *      avl_rotate(root_ptr, A, +sign);
 230  *
 231  */
 232 static AVL_INLINE struct avl_tree_node *
 233 avl_do_double_rotate(struct avl_tree_node ** const root_ptr,
 234                      struct avl_tree_node * const B,
 235                      struct avl_tree_node * const A, const int sign)
 236 {
 237         struct avl_tree_node * const E = avl_get_child(B, +sign);
 238         struct avl_tree_node * const F = avl_get_child(E, -sign);
 239         struct avl_tree_node * const G = avl_get_child(E, +sign);
 240         struct avl_tree_node * const X = avl_get_parent(A);
 241         const int e = avl_get_balance_factor(E);
 242
 243         avl_set_child(A, -sign, G);
 244         avl_set_parent_balance(A, E, ((sign * e >= 0) ? 0 : -e));
 245
 246         avl_set_child(B, +sign, F);
 247         avl_set_parent_balance(B, E, ((sign * e <= 0) ? 0 : -e));
 248
 249         avl_set_child(E, +sign, A);
 250         avl_set_child(E, -sign, B);
 251         avl_set_parent_balance(E, X, 0);
 252
 253         if (G)
 254                 avl_set_parent(G, A);
 255
 256         if (F)
 257                 avl_set_parent(F, B);
 258
 259         if (X) {
 260                 if (avl_get_child(X, +sign) == A)
 261                         avl_set_child(X, +sign, E);
 262                 else
 263                         avl_set_child(X, -sign, E);
 264         } else {
 265                 *root_ptr = E;
 266         }
 267
 268         return E;
 269 }
 270
 271 static AVL_INLINE bool
 272 avl_handle_subtree_growth(struct avl_tree_node ** const root_ptr,
 273                           struct avl_tree_node * const node,
 274                           struct avl_tree_node * const parent,
 275                           const int sign)
 276 {
 277         /* The subtree rooted at @node, which is a child of @parent, has
 278          * increased by 1.
 279          *
 280          * Adjust @parent's balance factor and check whether rotations need to
 281          * be done.  */
 282
 283         int old_balance_factor, new_balance_factor;
 284
 285         old_balance_factor = avl_get_balance_factor(parent);
 286
 287         if (old_balance_factor == 0) {
 288                 /* Parent has increased in height but is still sufficiently
 289                  * balanced.  Continue up the tree.  */
 290                 avl_adjust_balance_factor(parent, sign);
 291                 return false;
 292         }
 293
 294         new_balance_factor = old_balance_factor + sign;
 295
 296         if (new_balance_factor == 0) {
 297                 /* Parent is balanced; nothing more to to.  */
 298                 avl_adjust_balance_factor(parent, sign);
 299                 return true;
 300         }
 301
 302         /* FROM THIS POINT ONWARDS THE COMMENTS ASSUME sign < 0.
 303          * The other case is symmetric --- that is, the rotations done are the
 304          * the mirror images, all the balance factors are inverted, and left and
 305          * right pointers are otherwise reversed.  */
 306
 307         /* Parent is left-heavy (balance_factor == -2).  */
 308
 309         if (sign * avl_get_balance_factor(node) > 0) {
 310
 311                 /* Child node (@node, also B below) is also left-heavy.
 312                  * It must have balance_factor == -1.
 313                  * Do a clockwise ("right") rotation rooted at
 314                  * @parent (A below):
 315                  *
 316                  *           A              B
 317                  *          / \           /   \
 318                  *         B   C   =>    D     A
 319                  *        / \           / \   / \
 320                  *       D   E         F   G E   C
 321                  *      / \
 322                  *     F   G
 323                  *
 324                  * Before the rotation:
 325                  *      balance(A) = -2
 326                  *      balance(B) = -1
 327                  * Let x = height(C).  Then:
 328                  *      height(B) = x + 2
 329                  *      height(D) = x + 1
 330                  *      height(E) = x
 331                  *      max(height(F), height(G)) = x.
 332                  *
 333                  * After the rotation:
 334                  *      height(D) = max(height(F), height(G)) + 1
 335                  *                = x + 1
 336                  *      height(A) = max(height(E), height(C)) + 1
 337                  *                = max(x, x) + 1 = x + 1
 338                  *      balance(B) = 0
 339                  *      balance(A) = 0
 340                  */
 341
 342                 avl_rotate(root_ptr, parent, -sign);
 343                 avl_adjust_balance_factor(parent, -sign); /* A */
 344                 avl_adjust_balance_factor(node, -sign);   /* B */
 345         } else {
 346                 /* Child node (@node, also B below) is right-heavy.
 347                  * It must have balance_factor == +1.
 348                  * Do a counterclockwise ("left") rotation rooted at child node
 349                  * (B below), then a clockwise ("right") rotation rooted at
 350                  * parent node (A below).
 351                  *
 352                  *           A             A           E
 353                  *          / \           / \        /   \
 354                  *         B   C   =>    E   C  =>  B     A
 355                  *        / \           / \        / \   / \
 356                  *       D   E         B   G      D   F G   C
 357                  *          / \       / \
 358                  *         F   G     D   F
 359                  *
 360                  * Before the rotation:
 361                  *      balance(A) = -2
 362                  *      balance(B) = +1
 363                  * Let x = height(C).  Then:
 364                  *      height(B) = x + 2
 365                  *      height(E) = x + 1
 366                  *      height(D) = x
 367                  *      max(height(F), height(G)) = x
 368                  *
 369                  * After both rotations:
 370                  *      height(A) = max(height(G), height(C)) + 1
 371                  *                = x + 1
 372                  *      balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
 373                  *      height(B) = max(height(D), height(F)) + 1
 374                  *                = x + 1
 375                  *      balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
 376                  *
 377                  *      height(E) = x + 2
 378                  *      balance(E) = 0
 379                  */
 380                 avl_do_double_rotate(root_ptr, node, parent, -sign);
 381         }
 382         return true;
 383 }
 384
 385 /* Rebalance the tree after insertion of the specified node.  */
 386 void
 387 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
 388                                 struct avl_tree_node *inserted)
 389 {
 390         struct avl_tree_node *node, *parent;
 391         bool done;
 392
 393         inserted->left = NULL;
 394         inserted->right = NULL;
 395
 396         node = inserted;
 397
 398         /* Adjust balance factor of new node's parent.
 399          * No rotation will need to be done at this level.  */
 400
 401         parent = avl_get_parent(node);
 402         if (!parent)
 403                 return;
 404
 405         if (node == parent->left)
 406                 avl_adjust_balance_factor(parent, -1);
 407         else
 408                 avl_adjust_balance_factor(parent, +1);
 409
 410         if (avl_get_balance_factor(parent) == 0)
 411                 /* @parent did not change in height.  Nothing more to do.  */
 412                 return;
 413
 414         /* The subtree rooted at @parent increased in height by 1.  */
 415
 416         do {
 417                 /* Adjust balance factor of next ancestor.  */
 418
 419                 node = parent;
 420                 parent = avl_get_parent(node);
 421                 if (!parent)
 422                         return;
 423
 424                 /* The subtree rooted at @node has increased in height by 1.  */
 425                 if (node == parent->left)
 426                         done = avl_handle_subtree_growth(root_ptr, node,
 427                                                          parent, -1);
 428                 else
 429                         done = avl_handle_subtree_growth(root_ptr, node,
 430                                                          parent, +1);
 431         } while (!done);
 432 }
 433
 434 static AVL_INLINE struct avl_tree_node *
 435 avl_handle_subtree_shrink(struct avl_tree_node ** const root_ptr,
 436                           struct avl_tree_node *parent,
 437                           const int sign,
 438                           bool * const left_deleted_ret)
 439 {
 440         struct avl_tree_node *node;
 441
 442         const int old_balance_factor = avl_get_balance_factor(parent);
 443         const int new_balance_factor = old_balance_factor + sign;
 444
 445         if (old_balance_factor == 0) {
 446                 /* Prior to the deletion, the subtree rooted at
 447                  * @parent was perfectly balanced.  It's now
 448                  * unbalanced by 1, but that's okay and its height
 449                  * hasn't changed.  Nothing more to do.  */
 450                 avl_adjust_balance_factor(parent, sign);
 451                 return NULL;
 452         } else if (new_balance_factor == 0) {
 453                 /* The subtree rooted at @parent is now perfectly
 454                  * balanced, whereas before the deletion it was
 455                  * unbalanced by 1.  Its height must have decreased
 456                  * by 1.  No rotation is needed at this location,
 457                  * but continue up the tree.  */
 458                 avl_adjust_balance_factor(parent, sign);
 459                 node = parent;
 460         } else {
 461                 /* The subtree rooted at @parent is now significantly
 462                  * unbalanced (by 2 in some direction).  */
 463                 node = avl_get_child(parent, sign);
 464
 465                 /* The rotations below are similar to those done during
 466                  * insertion.  The only new case is the one where the
 467                  * child node has a balance factor of 0.  */
 468
 469                 if (sign * avl_get_balance_factor(node) >= 0) {
 470                         avl_rotate(root_ptr, parent, -sign);
 471
 472                         if (avl_get_balance_factor(node) == 0) {
 473                                 /* Child node (@node, also B below) is balanced.
 474                                  * Do a clockwise ("right") rotation rooted at
 475                                  * @parent (A below):
 476                                  *
 477                                  *           A              B
 478                                  *          / \           /   \
 479                                  *         B   C   =>    D     A
 480                                  *        / \           / \   / \
 481                                  *       D   E         F   G E   C
 482                                  *      / \
 483                                  *     F   G
 484                                  *
 485                                  * Before the rotation:
 486                                  *      balance(A) = -2
 487                                  *      balance(B) =  0
 488                                  * Let x = height(C).  Then:
 489                                  *      height(B) = x + 2
 490                                  *      height(D) = x + 1
 491                                  *      height(E) = x + 1
 492                                  *      max(height(F), height(G)) = x.
 493                                  *
 494                                  * After the rotation:
 495                                  *      height(D) = max(height(F), height(G)) + 1
 496                                  *                = x + 1
 497                                  *      height(A) = max(height(E), height(C)) + 1
 498                                  *                = max(x + 1, x) + 1 = x + 2
 499                                  *      balance(A) = -1
 500                                  *      balance(B) = +1
 501                                  */
 502
 503                                 /* A: -2 => -1 (sign < 0)
 504                                  * or +2 => +1 (sign > 0)
 505                                  * No change needed --- that's the same as
 506                                  * old_balance_factor.  */
 507
 508                                 /* B: 0 => +1 (sign < 0)
 509                                  * or 0 => -1 (sign > 0)  */
 510                                 avl_adjust_balance_factor(node, -sign);
 511
 512                                 /* Height is unchanged; nothing more to do.  */
 513                                 return NULL;
 514                         } else {
 515                                 avl_adjust_balance_factor(parent, -sign);
 516                                 avl_adjust_balance_factor(node, -sign);
 517                         }
 518                 } else {
 519                         node = avl_do_double_rotate(root_ptr, node,
 520                                                     parent, -sign);
 521                 }
 522         }
 523         parent = avl_get_parent(node);
 524         if (parent)
 525                 *left_deleted_ret = (node == parent->left);
 526         return parent;
 527 }
 528
 529 /* Swaps node X, which must have 2 children, with its in-order successor, then
 530  * unlinks node X.  Returns the parent of X just before unlinking, without its
 531  * balance factor having been updated to account for the unlink.  */
 532 static AVL_INLINE struct avl_tree_node *
 533 avl_tree_swap_with_successor(struct avl_tree_node **root_ptr,
 534                              struct avl_tree_node *X,
 535                              bool *left_deleted_ret)
 536 {
 537         struct avl_tree_node *Y, *P, *ret;
 538
 539         Y = X->right;
 540         if (!Y->left) {
 541                 /*
 542                  *     P?           P?           P?
 543                  *     |            |            |
 544                  *     X            Y            Y
 545                  *    / \          / \          / \
 546                  *   A   Y    =>  A   X    =>  A   B?
 547                  *      / \          / \
 548                  *    (0)  B?      (0)  B?
 549                  *
 550                  * [ X removed, Y returned ]
 551                  */
 552                 ret = Y;
 553                 *left_deleted_ret = false;
 554         } else {
 555                 struct avl_tree_node *Q;
 556
 557                 do {
 558                         Q = Y;
 559                         Y = Y->left;
 560                 } while (Y->left);
 561
 562                 /*
 563                  *     P?           P?           P?
 564                  *     |            |            |
 565                  *     X            Y            Y
 566                  *    / \          / \          / \
 567                  *   A   ...  =>  A  ...   =>  A  ...
 568                  *       |            |            |
 569                  *       Q            Q            Q
 570                  *      /            /            /
 571                  *     Y            X            B?
 572                  *    / \          / \
 573                  *  (0)  B?      (0)  B?
 574                  *
 575                  *
 576                  * [ X removed, Q returned ]
 577                  */
 578
 579                 Q->left = Y->right;
 580                 if (Q->left)
 581                         avl_set_parent(Q->left, Q);
 582                 Y->right = X->right;
 583                 avl_set_parent(X->right, Y);
 584                 ret = Q;
 585                 *left_deleted_ret = true;
 586         }
 587
 588         Y->left = X->left;
 589         avl_set_parent(X->left, Y);
 590
 591         Y->parent_balance = X->parent_balance;
 592         P = avl_get_parent(X);
 593         if (P) {
 594                 if (P->left == X)
 595                         P->left = Y;
 596                 else
 597                         P->right = Y;
 598         } else {
 599                 *root_ptr = Y;
 600         }
 601
 602         return ret;
 603 }
 604
 605 /* Removes the specified @node from the AVL tree.  @root_ptr must point to the
 606  * pointer to the root node of the tree; *root_ptr may change if the tree is
 607  * rebalanced.
 608  *
 609  * This *only* removes the node and rebalances the tree; it does not free
 610  * memory, nor does it do the equivalent of avl_tree_node_set_unlinked().  */
 611 void
 612 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
 613 {
 614         struct avl_tree_node *child, *parent;
 615         bool left_deleted;
 616
 617         if (node->left && node->right) {
 618                 parent = avl_tree_swap_with_successor(root_ptr, node,
 619                                                       &left_deleted);
 620         } else {
 621                 /* Unlink @node  */
 622                 child = node->left ? node->left : node->right;
 623                 parent = avl_get_parent(node);
 624                 if (parent) {
 625                         if (node == parent->left) {
 626                                 parent->left = child;
 627                                 left_deleted = true;
 628                         } else {
 629                                 parent->right = child;
 630                                 left_deleted = false;
 631                         }
 632                 } else {
 633                         *root_ptr = child;
 634                 }
 635                 if (child)
 636                         avl_set_parent(child, parent);
 637                 if (!parent)
 638                         return;
 639         }
 640
 641         /* Rebalance the tree  */
 642         do {
 643                 if (left_deleted)
 644                         parent = avl_handle_subtree_shrink(root_ptr, parent,
 645                                                            +1, &left_deleted);
 646                 else
 647                         parent = avl_handle_subtree_shrink(root_ptr, parent,
 648                                                            -1, &left_deleted);
 649         } while (parent);
 650 }