4 * Intrusive, nonrecursive AVL tree data structure (self-balancing binary search
5 * tree), implementation file.
10 * The author dedicates this file to the public domain.
11 * You can do whatever you want with this file.
14 #include <wimlib/avl_tree.h>
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)
21 const struct avl_tree_node *first = root;
26 return (struct avl_tree_node *)first;
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)
34 const struct avl_tree_node *next;
37 for (next = prev->right;
42 for (next = avl_get_parent(prev);
43 next && prev == next->right;
44 prev = next, next = avl_get_parent(next))
46 return (struct avl_tree_node *)next;
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)
53 const struct avl_tree_node *first = root;
56 while (first->left || first->right)
57 first = first->left ? first->left : first->right;
59 return (struct avl_tree_node *)first;
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)
70 const struct avl_tree_node *next = prev_parent;
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)
77 return (struct avl_tree_node *)next;
80 /* Returns the left child (sign < 0) or the right child (sign > 0) of the
81 * specified AVL tree node.
82 * Note: for all calls of this, 'sign' is constant at compilation time,
83 * so the compiler can remove the conditional. */
84 static AVL_INLINE struct avl_tree_node *
85 avl_get_child(const struct avl_tree_node *parent, int sign)
93 /* Sets the left child (sign < 0) or the right child (sign > 0) of the
94 * specified AVL tree node.
95 * Note: for all calls of this, 'sign' is constant at compilation time,
96 * so the compiler can remove the conditional. */
97 static AVL_INLINE void
98 avl_set_child(struct avl_tree_node *parent, int sign,
99 struct avl_tree_node *child)
102 parent->left = child;
104 parent->right = child;
107 /* Sets the parent and balance factor of the specified AVL tree node. */
108 static AVL_INLINE void
109 avl_set_parent_balance(struct avl_tree_node *node, struct avl_tree_node *parent,
112 node->parent_balance = (uintptr_t)parent | (balance_factor + 1);
115 /* Sets the parent of the specified AVL tree node. */
116 static AVL_INLINE void
117 avl_set_parent(struct avl_tree_node *node, struct avl_tree_node *parent)
119 node->parent_balance = (uintptr_t)parent | (node->parent_balance & 3);
122 /* Returns the balance factor of the specified AVL tree node --- that is, the
123 * height of its right subtree minus the height of its left subtree. */
124 static AVL_INLINE int
125 avl_get_balance_factor(const struct avl_tree_node *node)
127 return (int)(node->parent_balance & 3) - 1;
130 /* Sets the balance factor of the specified AVL tree node. This must be
132 static AVL_INLINE void
133 avl_set_balance_factor(struct avl_tree_node *node, int balance_factor)
135 node->parent_balance =
136 (node->parent_balance & ~3) | (balance_factor + 1);
139 /* Adds @amount to the balance factor of the specified AVL tree node.
140 * The caller must ensure this still results in a valid balance factor
142 static AVL_INLINE void
143 avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
145 node->parent_balance += amount;
148 static AVL_INLINE void
149 avl_replace_child(struct avl_tree_node **root_ptr,
150 struct avl_tree_node *parent,
151 struct avl_tree_node *old_child,
152 struct avl_tree_node *new_child)
155 if (old_child == parent->left)
156 parent->left = new_child;
158 parent->right = new_child;
160 *root_ptr = new_child;
165 * Template for performing a single rotation ---
167 * sign > 0: Rotate clockwise (right) rooted at A:
177 * (nodes marked with ? may not exist)
179 * sign < 0: Rotate counterclockwise (left) rooted at A:
189 * This updates pointers but not balance factors!
191 static AVL_INLINE void
192 avl_rotate(struct avl_tree_node ** const root_ptr,
193 struct avl_tree_node * const A, const int sign)
195 struct avl_tree_node * const B = avl_get_child(A, -sign);
196 struct avl_tree_node * const E = avl_get_child(B, +sign);
197 struct avl_tree_node * const P = avl_get_parent(A);
199 avl_set_child(A, -sign, E);
200 avl_set_parent(A, B);
202 avl_set_child(B, +sign, A);
203 avl_set_parent(B, P);
206 avl_set_parent(E, A);
208 avl_replace_child(root_ptr, P, A, B);
212 * Template for performing a double rotation ---
214 * sign > 0: Rotate counterclockwise (left) rooted at B, then
215 * clockwise (right) rooted at A:
221 * B C? => E C? => B A
223 * D? E B G? D? F?G? C?
227 * (nodes marked with ? may not exist)
229 * sign < 0: Rotate clockwise (right) rooted at B, then
230 * counterclockwise (left) rooted at A:
236 * C? B => C? E => A B
238 * E D? G? B C? G?F? D?
242 * Returns a pointer to E and updates balance factors. Except for those
243 * two things, this function is equivalent to:
244 * avl_rotate(root_ptr, B, -sign);
245 * avl_rotate(root_ptr, A, +sign);
247 * See comment in avl_handle_subtree_growth() for explanation of balance
250 static AVL_INLINE struct avl_tree_node *
251 avl_do_double_rotate(struct avl_tree_node ** const root_ptr,
252 struct avl_tree_node * const B,
253 struct avl_tree_node * const A, const int sign)
255 struct avl_tree_node * const E = avl_get_child(B, +sign);
256 struct avl_tree_node * const F = avl_get_child(E, -sign);
257 struct avl_tree_node * const G = avl_get_child(E, +sign);
258 struct avl_tree_node * const P = avl_get_parent(A);
259 const int e = avl_get_balance_factor(E);
261 avl_set_child(A, -sign, G);
262 avl_set_parent_balance(A, E, ((sign * e >= 0) ? 0 : -e));
264 avl_set_child(B, +sign, F);
265 avl_set_parent_balance(B, E, ((sign * e <= 0) ? 0 : -e));
267 avl_set_child(E, +sign, A);
268 avl_set_child(E, -sign, B);
269 avl_set_parent_balance(E, P, 0);
272 avl_set_parent(G, A);
275 avl_set_parent(F, B);
277 avl_replace_child(root_ptr, P, A, E);
283 * This function handles the growth of a subtree due to an insertion.
286 * Location of the tree's root pointer.
289 * A subtree that has increased in height by 1 due to an insertion.
292 * Parent of @node; must not be NULL.
295 * -1 if @node is the left child of @parent;
296 * +1 if @node is the right child of @parent.
298 * This function will adjust @parent's balance factor, then do a (single
299 * or double) rotation if necessary. The return value will be %true if
300 * the full AVL tree is now adequately balanced, or %false if the subtree
301 * rooted at @parent is now adequately balanced but has increased in
302 * height by 1, so the caller should continue up the tree.
304 * Note that if %false is returned, no rotation will have been done.
305 * Indeed, a single node insertion cannot require that more than one
306 * (single or double) rotation be done.
308 static AVL_INLINE bool
309 avl_handle_subtree_growth(struct avl_tree_node ** const root_ptr,
310 struct avl_tree_node * const node,
311 struct avl_tree_node * const parent,
314 int old_balance_factor, new_balance_factor;
316 old_balance_factor = avl_get_balance_factor(parent);
318 if (old_balance_factor == 0) {
319 avl_adjust_balance_factor(parent, sign);
320 /* @parent is still sufficiently balanced (-1 or +1
321 * balance factor), but must have increased in height.
322 * Continue up the tree. */
326 new_balance_factor = old_balance_factor + sign;
328 if (new_balance_factor == 0) {
329 avl_adjust_balance_factor(parent, sign);
330 /* @parent is now perfectly balanced (0 balance factor).
331 * It cannot have increased in height, so there is
332 * nothing more to do. */
336 /* @parent is too left-heavy (new_balance_factor == -2) or
337 * too right-heavy (new_balance_factor == +2). */
339 /* Test whether @node is left-heavy (-1 balance factor) or
340 * right-heavy (+1 balance factor).
341 * Note that it cannot be perfectly balanced (0 balance factor)
342 * because here we are under the invariant that @node has
343 * increased in height due to the insertion. */
344 if (sign * avl_get_balance_factor(node) > 0) {
346 /* @node (B below) is heavy in the same direction @parent
347 * (A below) is heavy.
349 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
350 * The comment, diagram, and equations below assume sign < 0.
351 * The other case is symmetric!
352 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
354 * Do a clockwise rotation rooted at @parent (A below):
364 * Before the rotation:
367 * Let x = height(C). Then:
371 * max(height(F), height(G)) = x.
373 * After the rotation:
374 * height(D) = max(height(F), height(G)) + 1
376 * height(A) = max(height(E), height(C)) + 1
377 * = max(x, x) + 1 = x + 1
381 avl_rotate(root_ptr, parent, -sign);
384 * avl_set_balance_factor(parent, 0); */
385 avl_adjust_balance_factor(parent, -sign); /* A */
388 * avl_set_balance_factor(node, 0); */
389 avl_adjust_balance_factor(node, -sign); /* B */
391 /* @node (B below) is heavy in the direction opposite
392 * from the direction @parent (A below) is heavy.
394 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
395 * The comment, diagram, and equations below assume sign < 0.
396 * The other case is symmetric!
397 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
399 * Do a counterblockwise rotation rooted at @node (B below),
400 * then a clockwise rotation rooted at @parent (A below):
404 * B C? => E C? => B A
406 * D? E B G? D? F?G? C?
410 * Before the rotation:
413 * Let x = height(C). Then:
417 * max(height(F), height(G)) = x
419 * After both rotations:
420 * height(A) = max(height(G), height(C)) + 1
422 * balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
423 * height(B) = max(height(D), height(F)) + 1
425 * balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
430 avl_do_double_rotate(root_ptr, node, parent, -sign);
433 /* Height after rotation is unchanged; nothing more to do. */
437 /* Rebalance the tree after insertion of the specified node. */
439 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
440 struct avl_tree_node *inserted)
442 struct avl_tree_node *node, *parent;
445 inserted->left = NULL;
446 inserted->right = NULL;
450 /* Adjust balance factor of new node's parent.
451 * No rotation will need to be done at this level. */
453 parent = avl_get_parent(node);
457 if (node == parent->left)
458 avl_adjust_balance_factor(parent, -1);
460 avl_adjust_balance_factor(parent, +1);
462 if (avl_get_balance_factor(parent) == 0)
463 /* @parent did not change in height. Nothing more to do. */
466 /* The subtree rooted at @parent increased in height by 1. */
469 /* Adjust balance factor of next ancestor. */
472 parent = avl_get_parent(node);
476 /* The subtree rooted at @node has increased in height by 1. */
477 if (node == parent->left)
478 done = avl_handle_subtree_growth(root_ptr, node,
481 done = avl_handle_subtree_growth(root_ptr, node,
487 * This function handles the shrinkage of a subtree due to a deletion.
490 * Location of the tree's root pointer.
493 * A node in the tree, exactly one of whose subtrees has decreased
494 * in height by 1 due to a deletion. (This includes the case where
495 * one of the child pointers has become NULL, since we can consider
496 * the "NULL" subtree to have a height of 0.)
499 * +1 if the left subtree of @parent has decreased in height by 1;
500 * -1 if the right subtree of @parent has decreased in height by 1.
503 * If the return value is not NULL, this will be set to %true if the
504 * left subtree of the returned node has decreased in height by 1,
505 * or %false if the right subtree of the returned node has decreased
508 * This function will adjust @parent's balance factor, then do a (single
509 * or double) rotation if necessary. The return value will be NULL if
510 * the full AVL tree is now adequately balanced, or a pointer to the
511 * parent of @parent if @parent is now adequately balanced but has
512 * decreased in height by 1. Also in the latter case, *left_deleted_ret
515 static AVL_INLINE struct avl_tree_node *
516 avl_handle_subtree_shrink(struct avl_tree_node ** const root_ptr,
517 struct avl_tree_node *parent,
519 bool * const left_deleted_ret)
521 struct avl_tree_node *node;
522 int old_balance_factor, new_balance_factor;
524 old_balance_factor = avl_get_balance_factor(parent);
526 if (old_balance_factor == 0) {
527 /* Prior to the deletion, the subtree rooted at
528 * @parent was perfectly balanced. It's now
529 * unbalanced by 1, but that's okay and its height
530 * hasn't changed. Nothing more to do. */
531 avl_adjust_balance_factor(parent, sign);
535 new_balance_factor = old_balance_factor + sign;
537 if (new_balance_factor == 0) {
538 /* The subtree rooted at @parent is now perfectly
539 * balanced, whereas before the deletion it was
540 * unbalanced by 1. Its height must have decreased
541 * by 1. No rotation is needed at this location,
542 * but continue up the tree. */
543 avl_adjust_balance_factor(parent, sign);
546 /* @parent is too left-heavy (new_balance_factor == -2) or
547 * too right-heavy (new_balance_factor == +2). */
549 node = avl_get_child(parent, sign);
551 /* The rotations below are similar to those done during
552 * insertion (see avl_handle_subtree_growth()), so full
553 * comments are not provided. The only new case is the
554 * one where @node has a balance factor of 0, and that is
557 if (sign * avl_get_balance_factor(node) >= 0) {
559 avl_rotate(root_ptr, parent, -sign);
561 if (avl_get_balance_factor(node) == 0) {
563 * @node (B below) is perfectly balanced.
565 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
566 * The comment, diagram, and equations
567 * below assume sign < 0. The other case
569 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
571 * Do a clockwise rotation rooted at
582 * Before the rotation:
585 * Let x = height(C). Then:
589 * max(height(F), height(G)) = x.
591 * After the rotation:
592 * height(D) = max(height(F), height(G)) + 1
594 * height(A) = max(height(E), height(C)) + 1
595 * = max(x + 1, x) + 1 = x + 2
600 /* A: -2 => -1 (sign < 0)
601 * or +2 => +1 (sign > 0)
602 * No change needed --- that's the same as
603 * old_balance_factor. */
605 /* B: 0 => +1 (sign < 0)
606 * or 0 => -1 (sign > 0) */
607 avl_adjust_balance_factor(node, -sign);
609 /* Height is unchanged; nothing more to do. */
612 avl_adjust_balance_factor(parent, -sign);
613 avl_adjust_balance_factor(node, -sign);
616 node = avl_do_double_rotate(root_ptr, node,
620 parent = avl_get_parent(node);
622 *left_deleted_ret = (node == parent->left);
626 /* Swaps node X, which must have 2 children, with its in-order successor, then
627 * unlinks node X. Returns the parent of X just before unlinking, without its
628 * balance factor having been updated to account for the unlink. */
629 static AVL_INLINE struct avl_tree_node *
630 avl_tree_swap_with_successor(struct avl_tree_node **root_ptr,
631 struct avl_tree_node *X,
632 bool *left_deleted_ret)
634 struct avl_tree_node *Y, *ret;
647 * [ X unlinked, Y returned ]
650 *left_deleted_ret = false;
652 struct avl_tree_node *Q;
664 * A ... => A ... => A ...
673 * [ X unlinked, Q returned ]
678 avl_set_parent(Q->left, Q);
680 avl_set_parent(X->right, Y);
682 *left_deleted_ret = true;
686 avl_set_parent(X->left, Y);
688 Y->parent_balance = X->parent_balance;
689 avl_replace_child(root_ptr, avl_get_parent(X), X, Y);
695 * Removes an item from the specified AVL tree.
698 * Location of the AVL tree's root pointer. Indirection is needed
699 * because the root node may change if the tree needed to be rebalanced
700 * because of the deletion or if @node was the root node.
703 * Pointer to the `struct avl_tree_node' embedded in the item to
704 * remove from the tree.
706 * Note: This function *only* removes the node and rebalances the tree.
707 * It does not free any memory, nor does it do the equivalent of
708 * avl_tree_node_set_unlinked().
711 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
713 struct avl_tree_node *parent;
714 bool left_deleted = false;
716 if (node->left && node->right) {
717 /* @node is fully internal, with two children. Swap it
718 * with its in-order successor (which must exist in the
719 * right subtree of @node and can have, at most, a right
720 * child), then unlink @node. */
721 parent = avl_tree_swap_with_successor(root_ptr, node,
723 /* @parent is now the parent of what was @node's in-order
724 * successor. It cannot be NULL, since @node itself was
725 * an ancestor of its in-order successor.
726 * @left_deleted has been set to %true if @node's
727 * in-order successor was the left child of @parent,
728 * otherwise %false. */
730 struct avl_tree_node *child;
732 /* @node is missing at least one child. Unlink it. Set
733 * @parent to @node's parent, and set @left_deleted to
734 * reflect which child of @parent @node was. Or, if
735 * @node was the root node, simply update the root node
737 child = node->left ? node->left : node->right;
738 parent = avl_get_parent(node);
740 if (node == parent->left) {
741 parent->left = child;
744 parent->right = child;
745 left_deleted = false;
748 avl_set_parent(child, parent);
751 avl_set_parent(child, parent);
757 /* Rebalance the tree. */
760 parent = avl_handle_subtree_shrink(root_ptr, parent,
763 parent = avl_handle_subtree_shrink(root_ptr, parent,