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 /* Returns the balance factor of the specified AVL tree node --- that is, the
 114  * height of its right subtree minus the height of its left subtree.  */
 115 static AVL_INLINE int
 116 avl_get_balance_factor(const struct avl_tree_node *node)
 117 {
 118         return (int)(node->parent_balance & 3) - 1;
 119 }
 120
 121 /* Sets the balance factor of the specified AVL tree node.  The caller MUST
 122  * ensure that this number is valid and maintains the AVL tree invariants.  */
 123 static AVL_INLINE void
 124 avl_set_balance_factor(struct avl_tree_node *node, int balance_factor)
 125 {
 126         node->parent_balance =
 127                 (node->parent_balance & ~3) | (balance_factor + 1);
 128 }
 129
 130 /* Increments the balance factor of the specified AVL tree @node by the
 131  * specified @amount.  The caller MUST ensure that this number is valid and
 132  * maintains the AVL tree invariants.  */
 133 static AVL_INLINE void
 134 avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
 135 {
 136         node->parent_balance += amount;
 137 }
 138
 139 /*
 140  * Template for performing a rotation ---
 141  *
 142  * Clockwise/Right (sign > 0):
 143  *
 144  *           X             X
 145  *           |             |
 146  *           A             B
 147  *          / \           / \
 148  *         B   C   =>    D   A
 149  *        / \               / \
 150  *       D   E             E   C
 151  *
 152  * Counterclockwise/Left (sign < 0)- same, except left and right are reversed:
 153  *
 154  *           X             X
 155  *           |             |
 156  *           A             B
 157  *          / \           / \
 158  *         C   B   =>    A   D
 159  *            / \       / \
 160  *           E   D     C   E
 161  *
 162  * This updates pointers but not balance factors!
 163  */
 164 static AVL_INLINE void
 165 avl_rotate(struct avl_tree_node * const A, const int sign)
 166 {
 167         struct avl_tree_node * const B = avl_get_child(A, -sign);
 168         struct avl_tree_node * const C = avl_get_child(A, +sign);
 169         struct avl_tree_node * const E = avl_get_child(B, +sign);
 170         struct avl_tree_node * const X = avl_get_parent(A);
 171
 172         avl_set_child(A, -sign, E);
 173         avl_set_child(A, +sign, C);
 174         avl_set_parent(A, B);
 175
 176         avl_set_child(B, +sign, A);
 177         avl_set_parent(B, X);
 178
 179         if (E)
 180                 avl_set_parent(E, A);
 181
 182         if (X) {
 183                 if (avl_get_child(X, +sign) == A)
 184                         avl_set_child(X, +sign, B);
 185                 else
 186                         avl_set_child(X, -sign, B);
 187         }
 188 }
 189
 190 /* See description in avl_handle_subtree_growth()  */
 191 static AVL_INLINE struct avl_tree_node *
 192 avl_do_double_rotate(struct avl_tree_node * const B,
 193                      struct avl_tree_node * const A, const int sign)
 194 {
 195
 196         struct avl_tree_node * const E = avl_get_child(B, -sign);
 197         const int e = avl_get_balance_factor(E);
 198
 199         avl_rotate(B, +sign);
 200         avl_rotate(A, -sign);
 201         avl_set_balance_factor(A, ((sign * e <= 0) ? 0 : -e));
 202         avl_set_balance_factor(B, ((sign * e >= 0) ? 0 : -e));
 203         avl_set_balance_factor(E, 0);
 204
 205         return E;
 206 }
 207
 208 static AVL_INLINE bool
 209 avl_handle_subtree_growth(struct avl_tree_node * const node,
 210                           struct avl_tree_node * const parent,
 211                           const int sign)
 212 {
 213         /* Height of @node subtree of @parent has increased by 1.
 214          * Adjust @parent's balance factor and check whether rotations need to
 215          * be done.  */
 216
 217         const int old_balance_factor = avl_get_balance_factor(parent);
 218         const int new_balance_factor = old_balance_factor + sign;
 219
 220         if (old_balance_factor == 0) {
 221                 /* Parent has increased in height but is still sufficiently
 222                  * balanced.  Continue up the tree.  */
 223                 avl_adjust_balance_factor(parent, sign);
 224                 return false;
 225         }
 226
 227         if (new_balance_factor == 0) {
 228                 /* Parent is balanced; nothing more to to.  */
 229                 avl_adjust_balance_factor(parent, sign);
 230                 return true;
 231         }
 232
 233         /* FROM THIS POINT ONWARDS THE COMMENTS ASSUME sign < 0.
 234          * The other case is symmetric --- that is, the rotations done are the
 235          * the mirror images, all the balance factors are inverted, and left and
 236          * right pointers are otherwise reversed.  */
 237
 238         /* Parent is left-heavy (balance_factor == -2).  */
 239
 240         if (sign * avl_get_balance_factor(node) > 0) {
 241
 242                 /* Child node (@node, also B below) is also left-heavy.
 243                  * It must have balance_factor == -1.
 244                  * Do a clockwise ("right") rotation rooted at
 245                  * @parent (A below):
 246                  *
 247                  *           A              B
 248                  *          / \           /   \
 249                  *         B   C   =>    D     A
 250                  *        / \           / \   / \
 251                  *       D   E         F   G E   C
 252                  *      / \
 253                  *     F   G
 254                  *
 255                  * Before the rotation:
 256                  *      balance(A) = -2
 257                  *      balance(B) = -1
 258                  * Let x = height(C).  Then:
 259                  *      height(B) = x + 2
 260                  *      height(D) = x + 1
 261                  *      height(E) = x
 262                  *      max(height(F), height(G)) = x.
 263                  *
 264                  * After the rotation:
 265                  *      height(D) = max(height(F), height(G)) + 1
 266                  *                = x + 1
 267                  *      height(A) = max(height(E), height(C)) + 1
 268                  *                = max(x, x) + 1 = x + 1
 269                  *      balance(B) = 0
 270                  *      balance(A) = 0
 271                  */
 272
 273                 avl_rotate(parent, -sign);
 274
 275                 avl_set_balance_factor(node, 0);   /* B */
 276                 avl_set_balance_factor(parent, 0); /* A */
 277         } else {
 278                 /* Child node (@node, also B below) is right-heavy.
 279                  * It must have balance_factor == +1.
 280                  * Do a counterclockwise ("left") rotation rooted at child node
 281                  * (B below), then a clockwise ("right") rotation rooted at
 282                  * parent node (A below).
 283                  *
 284                  *           A             A           E
 285                  *          / \           / \        /   \
 286                  *         B   C   =>    E   C  =>  B     A
 287                  *        / \           / \        / \   / \
 288                  *       D   E         B   G      D   F G   C
 289                  *          / \       / \
 290                  *         F   G     D   F
 291                  *
 292                  * Before the rotation:
 293                  *      balance(A) = -2
 294                  *      balance(B) = +1
 295                  * Let x = height(C).  Then:
 296                  *      height(B) = x + 2
 297                  *      height(E) = x + 1
 298                  *      height(D) = x
 299                  *      max(height(F), height(G)) = x
 300                  *
 301                  * After both rotations:
 302                  *      height(A) = max(height(G), height(C)) + 1
 303                  *                = x + 1
 304                  *      balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
 305                  *      height(B) = max(height(D), height(F)) + 1
 306                  *                = x + 1
 307                  *      balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
 308                  *
 309                  *      height(E) = x + 2
 310                  *      balance(E) = 0
 311                  */
 312                 avl_do_double_rotate(node, parent, sign);
 313         }
 314         return true;
 315 }
 316
 317 /* Rebalance the tree after insertion of the specified node.  */
 318 void
 319 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
 320                                 struct avl_tree_node *inserted)
 321 {
 322         struct avl_tree_node *node, *parent;
 323         bool done = false;
 324
 325         inserted->left = NULL;
 326         inserted->right = NULL;
 327
 328         for (node = inserted, parent = avl_get_parent(node);
 329              parent && !done;
 330              node = parent, parent = avl_get_parent(parent))
 331         {
 332                 /* Height of @node subtree has increased by 1  */
 333
 334                 if (node == parent->left)
 335                         done = avl_handle_subtree_growth(node, parent, -1);
 336                 else
 337                         done = avl_handle_subtree_growth(node, parent, +1);
 338         }
 339         /* Due to rotations, *root_ptr may no longer be the root of the tree  */
 340         while (avl_get_parent(*root_ptr))
 341                 *root_ptr = avl_get_parent(*root_ptr);
 342 }
 343
 344 static AVL_INLINE struct avl_tree_node *
 345 avl_handle_subtree_shrink(struct avl_tree_node *parent,
 346                           const int sign,
 347                           bool * const left_deleted_ret)
 348 {
 349         struct avl_tree_node *node;
 350
 351         const int old_balance_factor = avl_get_balance_factor(parent);
 352         const int new_balance_factor = old_balance_factor + sign;
 353
 354         if (old_balance_factor == 0) {
 355                 /* Prior to the deletion, the subtree rooted at
 356                  * @parent was perfectly balanced.  It's now
 357                  * unbalanced by 1, but that's okay and its height
 358                  * hasn't changed.  Nothing more to do.  */
 359                 avl_adjust_balance_factor(parent, sign);
 360                 return NULL;
 361         } else if (new_balance_factor == 0) {
 362                 /* The subtree rooted at @parent is now perfectly
 363                  * balanced, whereas before the deletion it was
 364                  * unbalanced by 1.  Its height must have decreased
 365                  * by 1.  No rotation is needed at this location,
 366                  * but continue up the tree.  */
 367                 avl_adjust_balance_factor(parent, sign);
 368                 node = parent;
 369         } else {
 370                 /* The subtree rooted at @parent is now significantly
 371                  * unbalanced (by 2 in some direction).  */
 372                 node = avl_get_child(parent, sign);
 373
 374                 /* The rotations below are similar to those done during
 375                  * insertion.  The only new case is the one where the
 376                  * child node has a balance factor of 0.  */
 377
 378                 if (sign * avl_get_balance_factor(node) >= 0) {
 379                         avl_rotate(parent, -sign);
 380
 381                         if (avl_get_balance_factor(node) == 0) {
 382                                 avl_set_balance_factor(node,   -sign);
 383                                 avl_set_balance_factor(parent, +sign);
 384                                 /* Height is unchanged; nothing more to do.  */
 385                                 return NULL;
 386                         } else {
 387                                 avl_set_balance_factor(node, 0);
 388                                 avl_set_balance_factor(parent, 0);
 389                         }
 390                 } else {
 391                         node = avl_do_double_rotate(node, parent, sign);
 392                 }
 393         }
 394         parent = avl_get_parent(node);
 395         if (parent)
 396                 *left_deleted_ret = (node == parent->left);
 397         return parent;
 398 }
 399
 400 /* Swaps node X, which must have 2 children, with its in-order successor.
 401  * This adjusts all the necessary pointers and balance factors.  */
 402 static AVL_INLINE void
 403 avl_tree_swap_with_successor(struct avl_tree_node * const X)
 404 {
 405         struct avl_tree_node * const Y          = avl_tree_next_in_order(X);
 406         struct avl_tree_node * const Y_right    = Y->right;
 407         struct avl_tree_node * const Y_parent   = avl_get_parent(Y);
 408         struct avl_tree_node * const X_parent   = avl_get_parent(X);
 409         const int Y_balance_factor              = avl_get_balance_factor(Y);
 410
 411         /* Set child pointer to Y when placed in X's position  */
 412         if (X_parent) {
 413                 if (X == X_parent->left)
 414                         X_parent->left = Y;
 415                 else
 416                         X_parent->right = Y;
 417         }
 418
 419         /* Set child pointer to X when placed in Y's position.  Also set
 420          * the right pointer of Y (it may point to X if the nodes are
 421          * adjacent)  */
 422         if (Y_parent != X) {
 423                 if (Y == Y_parent->left)
 424                         Y_parent->left = X;
 425                 else
 426                         Y_parent->right = X;
 427                 Y->right = X->right;
 428         } else {
 429                 Y->right = X;
 430         }
 431
 432         /* Set parent pointer to X when placed in Y's position  */
 433         if (Y_right)
 434                 avl_set_parent(Y_right, X);
 435
 436         /* Note: Y has no left child since it is the in-order sucessor of X.  */
 437
 438         /* Set parent pointers to Y when placed in X's position, as well as X's
 439          * parent when placed in Y's position.  */
 440         avl_set_parent(X->left, Y);
 441         if (X->right != Y) {
 442                 avl_set_parent(X->right, Y);
 443                 avl_set_parent(X, Y_parent);
 444         } else {
 445                 avl_set_parent(X, Y);
 446         }
 447
 448         avl_set_parent(Y, X_parent);
 449         Y->left = X->left;
 450         avl_set_balance_factor(Y, avl_get_balance_factor(X));
 451
 452         X->left = NULL;
 453         X->right = Y_right;
 454         avl_set_balance_factor(X, Y_balance_factor);
 455 }
 456
 457 /* Removes the specified @node from the AVL tree.  @root_ptr must point to the
 458  * pointer to the root node of the tree; *root_ptr may change if the tree is
 459  * rebalanced.
 460  *
 461  * This *only* removes the node and rebalances the tree; it does not free
 462  * memory, nor does it do the equivalent of avl_tree_node_set_unlinked().  */
 463 void
 464 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
 465 {
 466         struct avl_tree_node *child, *parent;
 467         bool left_deleted;
 468
 469         if (node->left && node->right)
 470                 avl_tree_swap_with_successor(node);
 471
 472         /* Unlink @node  */
 473         child = node->left ? node->left : node->right;
 474         parent = avl_get_parent(node);
 475         if (parent) {
 476                 if (node == parent->left) {
 477                         parent->left = child;
 478                         left_deleted = true;
 479                 } else {
 480                         parent->right = child;
 481                         left_deleted = false;
 482                 }
 483         } else {
 484                 *root_ptr = child;
 485         }
 486         if (child)
 487                 avl_set_parent(child, parent);
 488
 489         /* Rebalance the tree  */
 490         while (parent) {
 491                 if (left_deleted)
 492                         parent = avl_handle_subtree_shrink(parent, +1, &left_deleted);
 493                 else
 494                         parent = avl_handle_subtree_shrink(parent, -1, &left_deleted);
 495         }
 496
 497         /* Due to rotations, *root_ptr may no longer point to the root of the
 498          * tree.  Fix it.  */
 499         if (*root_ptr)
 500                 while (avl_get_parent(*root_ptr))
 501                         *root_ptr = avl_get_parent(*root_ptr);
 502 }