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         /* Height of @node subtree of @parent has increased by 1.
 278          * Adjust @parent's balance factor and check whether rotations need to
 279          * be done.  */
 280
 281         int old_balance_factor, new_balance_factor;
 282
 283         old_balance_factor = avl_get_balance_factor(parent);
 284
 285         if (old_balance_factor == 0) {
 286                 /* Parent has increased in height but is still sufficiently
 287                  * balanced.  Continue up the tree.  */
 288                 avl_adjust_balance_factor(parent, sign);
 289                 return false;
 290         }
 291
 292         new_balance_factor = old_balance_factor + sign;
 293
 294         if (new_balance_factor == 0) {
 295                 /* Parent is balanced; nothing more to to.  */
 296                 avl_adjust_balance_factor(parent, sign);
 297                 return true;
 298         }
 299
 300         /* FROM THIS POINT ONWARDS THE COMMENTS ASSUME sign < 0.
 301          * The other case is symmetric --- that is, the rotations done are the
 302          * the mirror images, all the balance factors are inverted, and left and
 303          * right pointers are otherwise reversed.  */
 304
 305         /* Parent is left-heavy (balance_factor == -2).  */
 306
 307         if (sign * avl_get_balance_factor(node) > 0) {
 308
 309                 /* Child node (@node, also B below) is also left-heavy.
 310                  * It must have balance_factor == -1.
 311                  * Do a clockwise ("right") rotation rooted at
 312                  * @parent (A below):
 313                  *
 314                  *           A              B
 315                  *          / \           /   \
 316                  *         B   C   =>    D     A
 317                  *        / \           / \   / \
 318                  *       D   E         F   G E   C
 319                  *      / \
 320                  *     F   G
 321                  *
 322                  * Before the rotation:
 323                  *      balance(A) = -2
 324                  *      balance(B) = -1
 325                  * Let x = height(C).  Then:
 326                  *      height(B) = x + 2
 327                  *      height(D) = x + 1
 328                  *      height(E) = x
 329                  *      max(height(F), height(G)) = x.
 330                  *
 331                  * After the rotation:
 332                  *      height(D) = max(height(F), height(G)) + 1
 333                  *                = x + 1
 334                  *      height(A) = max(height(E), height(C)) + 1
 335                  *                = max(x, x) + 1 = x + 1
 336                  *      balance(B) = 0
 337                  *      balance(A) = 0
 338                  */
 339
 340                 avl_rotate(root_ptr, parent, -sign);
 341                 avl_adjust_balance_factor(parent, -sign); /* A */
 342                 avl_adjust_balance_factor(node, -sign);   /* B */
 343         } else {
 344                 /* Child node (@node, also B below) is right-heavy.
 345                  * It must have balance_factor == +1.
 346                  * Do a counterclockwise ("left") rotation rooted at child node
 347                  * (B below), then a clockwise ("right") rotation rooted at
 348                  * parent node (A below).
 349                  *
 350                  *           A             A           E
 351                  *          / \           / \        /   \
 352                  *         B   C   =>    E   C  =>  B     A
 353                  *        / \           / \        / \   / \
 354                  *       D   E         B   G      D   F G   C
 355                  *          / \       / \
 356                  *         F   G     D   F
 357                  *
 358                  * Before the rotation:
 359                  *      balance(A) = -2
 360                  *      balance(B) = +1
 361                  * Let x = height(C).  Then:
 362                  *      height(B) = x + 2
 363                  *      height(E) = x + 1
 364                  *      height(D) = x
 365                  *      max(height(F), height(G)) = x
 366                  *
 367                  * After both rotations:
 368                  *      height(A) = max(height(G), height(C)) + 1
 369                  *                = x + 1
 370                  *      balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
 371                  *      height(B) = max(height(D), height(F)) + 1
 372                  *                = x + 1
 373                  *      balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
 374                  *
 375                  *      height(E) = x + 2
 376                  *      balance(E) = 0
 377                  */
 378                 avl_do_double_rotate(root_ptr, node, parent, -sign);
 379         }
 380         return true;
 381 }
 382
 383 /* Rebalance the tree after insertion of the specified node.  */
 384 void
 385 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
 386                                 struct avl_tree_node *inserted)
 387 {
 388         struct avl_tree_node *node, *parent;
 389         bool done = false;
 390
 391         inserted->left = NULL;
 392         inserted->right = NULL;
 393
 394         for (node = inserted, parent = avl_get_parent(node);
 395              parent && !done;
 396              node = parent, parent = avl_get_parent(parent))
 397         {
 398                 /* Height of @node subtree has increased by 1  */
 399
 400                 if (node == parent->left)
 401                         done = avl_handle_subtree_growth(root_ptr, node,
 402                                                          parent, -1);
 403                 else
 404                         done = avl_handle_subtree_growth(root_ptr, node,
 405                                                          parent, +1);
 406         }
 407 }
 408
 409 static AVL_INLINE struct avl_tree_node *
 410 avl_handle_subtree_shrink(struct avl_tree_node ** const root_ptr,
 411                           struct avl_tree_node *parent,
 412                           const int sign,
 413                           bool * const left_deleted_ret)
 414 {
 415         struct avl_tree_node *node;
 416
 417         const int old_balance_factor = avl_get_balance_factor(parent);
 418         const int new_balance_factor = old_balance_factor + sign;
 419
 420         if (old_balance_factor == 0) {
 421                 /* Prior to the deletion, the subtree rooted at
 422                  * @parent was perfectly balanced.  It's now
 423                  * unbalanced by 1, but that's okay and its height
 424                  * hasn't changed.  Nothing more to do.  */
 425                 avl_adjust_balance_factor(parent, sign);
 426                 return NULL;
 427         } else if (new_balance_factor == 0) {
 428                 /* The subtree rooted at @parent is now perfectly
 429                  * balanced, whereas before the deletion it was
 430                  * unbalanced by 1.  Its height must have decreased
 431                  * by 1.  No rotation is needed at this location,
 432                  * but continue up the tree.  */
 433                 avl_adjust_balance_factor(parent, sign);
 434                 node = parent;
 435         } else {
 436                 /* The subtree rooted at @parent is now significantly
 437                  * unbalanced (by 2 in some direction).  */
 438                 node = avl_get_child(parent, sign);
 439
 440                 /* The rotations below are similar to those done during
 441                  * insertion.  The only new case is the one where the
 442                  * child node has a balance factor of 0.  */
 443
 444                 if (sign * avl_get_balance_factor(node) >= 0) {
 445                         avl_rotate(root_ptr, parent, -sign);
 446
 447                         if (avl_get_balance_factor(node) == 0) {
 448                                 /* Child node (@node, also B below) is balanced.
 449                                  * Do a clockwise ("right") rotation rooted at
 450                                  * @parent (A below):
 451                                  *
 452                                  *           A              B
 453                                  *          / \           /   \
 454                                  *         B   C   =>    D     A
 455                                  *        / \           / \   / \
 456                                  *       D   E         F   G E   C
 457                                  *      / \
 458                                  *     F   G
 459                                  *
 460                                  * Before the rotation:
 461                                  *      balance(A) = -2
 462                                  *      balance(B) =  0
 463                                  * Let x = height(C).  Then:
 464                                  *      height(B) = x + 2
 465                                  *      height(D) = x + 1
 466                                  *      height(E) = x + 1
 467                                  *      max(height(F), height(G)) = x.
 468                                  *
 469                                  * After the rotation:
 470                                  *      height(D) = max(height(F), height(G)) + 1
 471                                  *                = x + 1
 472                                  *      height(A) = max(height(E), height(C)) + 1
 473                                  *                = max(x + 1, x) + 1 = x + 2
 474                                  *      balance(A) = -1
 475                                  *      balance(B) = +1
 476                                  */
 477
 478                                 /* A: -2 => -1 (sign < 0)
 479                                  * or +2 => +1 (sign > 0)
 480                                  * No change needed --- that's the same as
 481                                  * old_balance_factor.  */
 482
 483                                 /* B: 0 => +1 (sign < 0)
 484                                  * or 0 => -1 (sign > 0)  */
 485                                 avl_adjust_balance_factor(node, -sign);
 486
 487                                 /* Height is unchanged; nothing more to do.  */
 488                                 return NULL;
 489                         } else {
 490                                 avl_adjust_balance_factor(parent, -sign);
 491                                 avl_adjust_balance_factor(node, -sign);
 492                         }
 493                 } else {
 494                         node = avl_do_double_rotate(root_ptr, node,
 495                                                     parent, -sign);
 496                 }
 497         }
 498         parent = avl_get_parent(node);
 499         if (parent)
 500                 *left_deleted_ret = (node == parent->left);
 501         return parent;
 502 }
 503
 504 /* Swaps node X, which must have 2 children, with its in-order successor, then
 505  * unlinks node X.  Returns the parent of X just before unlinking, without its
 506  * balance factor having been updated to account for the unlink.  */
 507 static AVL_INLINE struct avl_tree_node *
 508 avl_tree_swap_with_successor(struct avl_tree_node **root_ptr,
 509                              struct avl_tree_node *X,
 510                              bool *left_deleted_ret)
 511 {
 512         struct avl_tree_node *Y, *P, *ret;
 513
 514         Y = X->right;
 515         if (!Y->left) {
 516                 /*
 517                  *     P?           P?           P?
 518                  *     |            |            |
 519                  *     X            Y            Y
 520                  *    / \          / \          / \
 521                  *   A   Y    =>  A   X    =>  A   B?
 522                  *      / \          / \
 523                  *    (0)  B?      (0)  B?
 524                  *
 525                  * [ X removed, Y returned ]
 526                  */
 527                 ret = Y;
 528                 *left_deleted_ret = false;
 529         } else {
 530                 struct avl_tree_node *Q;
 531
 532                 do {
 533                         Q = Y;
 534                         Y = Y->left;
 535                 } while (Y->left);
 536
 537                 /*
 538                  *     P?           P?           P?
 539                  *     |            |            |
 540                  *     X            Y            Y
 541                  *    / \          / \          / \
 542                  *   A   ...  =>  A  ...   =>  A  ...
 543                  *       |            |            |
 544                  *       Q            Q            Q
 545                  *      /            /            /
 546                  *     Y            X            B?
 547                  *    / \          / \
 548                  *  (0)  B?      (0)  B?
 549                  *
 550                  *
 551                  * [ X removed, Q returned ]
 552                  */
 553
 554                 Q->left = Y->right;
 555                 if (Q->left)
 556                         avl_set_parent(Q->left, Q);
 557                 Y->right = X->right;
 558                 avl_set_parent(X->right, Y);
 559                 ret = Q;
 560                 *left_deleted_ret = true;
 561         }
 562
 563         Y->left = X->left;
 564         avl_set_parent(X->left, Y);
 565
 566         Y->parent_balance = X->parent_balance;
 567         P = avl_get_parent(X);
 568         if (P) {
 569                 if (P->left == X)
 570                         P->left = Y;
 571                 else
 572                         P->right = Y;
 573         } else {
 574                 *root_ptr = Y;
 575         }
 576
 577         return ret;
 578 }
 579
 580 /* Removes the specified @node from the AVL tree.  @root_ptr must point to the
 581  * pointer to the root node of the tree; *root_ptr may change if the tree is
 582  * rebalanced.
 583  *
 584  * This *only* removes the node and rebalances the tree; it does not free
 585  * memory, nor does it do the equivalent of avl_tree_node_set_unlinked().  */
 586 void
 587 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
 588 {
 589         struct avl_tree_node *child, *parent;
 590         bool left_deleted;
 591
 592         if (node->left && node->right) {
 593                 parent = avl_tree_swap_with_successor(root_ptr, node,
 594                                                       &left_deleted);
 595         } else {
 596                 /* Unlink @node  */
 597                 child = node->left ? node->left : node->right;
 598                 parent = avl_get_parent(node);
 599                 if (parent) {
 600                         if (node == parent->left) {
 601                                 parent->left = child;
 602                                 left_deleted = true;
 603                         } else {
 604                                 parent->right = child;
 605                                 left_deleted = false;
 606                         }
 607                 } else {
 608                         *root_ptr = child;
 609                 }
 610                 if (child)
 611                         avl_set_parent(child, parent);
 612                 if (!parent)
 613                         return;
 614         }
 615
 616         /* Rebalance the tree  */
 617         do {
 618                 if (left_deleted)
 619                         parent = avl_handle_subtree_shrink(root_ptr, parent,
 620                                                            +1, &left_deleted);
 621                 else
 622                         parent = avl_handle_subtree_shrink(root_ptr, parent,
 623                                                            -1, &left_deleted);
 624         } while (parent);
 625 }