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 /* Adds @amount to the balance factor of the specified AVL tree node.
131 * The caller must ensure this still results in a valid balance factor
133 static AVL_INLINE void
134 avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
136 node->parent_balance += amount;
139 static AVL_INLINE void
140 avl_replace_child(struct avl_tree_node **root_ptr,
141 struct avl_tree_node *parent,
142 struct avl_tree_node *old_child,
143 struct avl_tree_node *new_child)
146 if (old_child == parent->left)
147 parent->left = new_child;
149 parent->right = new_child;
151 *root_ptr = new_child;
156 * Template for performing a single rotation ---
158 * sign > 0: Rotate clockwise (right) rooted at A:
168 * (nodes marked with ? may not exist)
170 * sign < 0: Rotate counterclockwise (left) rooted at A:
180 * This updates pointers but not balance factors!
182 static AVL_INLINE void
183 avl_rotate(struct avl_tree_node ** const root_ptr,
184 struct avl_tree_node * const A, const int sign)
186 struct avl_tree_node * const B = avl_get_child(A, -sign);
187 struct avl_tree_node * const E = avl_get_child(B, +sign);
188 struct avl_tree_node * const P = avl_get_parent(A);
190 avl_set_child(A, -sign, E);
191 avl_set_parent(A, B);
193 avl_set_child(B, +sign, A);
194 avl_set_parent(B, P);
197 avl_set_parent(E, A);
199 avl_replace_child(root_ptr, P, A, B);
203 * Template for performing a double rotation ---
205 * sign > 0: Rotate counterclockwise (left) rooted at B, then
206 * clockwise (right) rooted at A:
212 * B C? => E C? => B A
214 * D? E B G? D? F?G? C?
218 * (nodes marked with ? may not exist)
220 * sign < 0: Rotate clockwise (right) rooted at B, then
221 * counterclockwise (left) rooted at A:
227 * C? B => C? E => A B
229 * E D? G? B C? G?F? D?
233 * Returns a pointer to E and updates balance factors. Except for those
234 * two things, this function is equivalent to:
235 * avl_rotate(root_ptr, B, -sign);
236 * avl_rotate(root_ptr, A, +sign);
238 * See comment in avl_handle_subtree_growth() for explanation of balance
241 static AVL_INLINE struct avl_tree_node *
242 avl_do_double_rotate(struct avl_tree_node ** const root_ptr,
243 struct avl_tree_node * const B,
244 struct avl_tree_node * const A, const int sign)
246 struct avl_tree_node * const E = avl_get_child(B, +sign);
247 struct avl_tree_node * const F = avl_get_child(E, -sign);
248 struct avl_tree_node * const G = avl_get_child(E, +sign);
249 struct avl_tree_node * const P = avl_get_parent(A);
250 const int e = avl_get_balance_factor(E);
252 avl_set_child(A, -sign, G);
253 avl_set_parent_balance(A, E, ((sign * e >= 0) ? 0 : -e));
255 avl_set_child(B, +sign, F);
256 avl_set_parent_balance(B, E, ((sign * e <= 0) ? 0 : -e));
258 avl_set_child(E, +sign, A);
259 avl_set_child(E, -sign, B);
260 avl_set_parent_balance(E, P, 0);
263 avl_set_parent(G, A);
266 avl_set_parent(F, B);
268 avl_replace_child(root_ptr, P, A, E);
274 * This function handles the growth of a subtree due to an insertion.
277 * Location of the tree's root pointer.
280 * A subtree that has increased in height by 1 due to an insertion.
283 * Parent of @node; must not be NULL.
286 * -1 if @node is the left child of @parent;
287 * +1 if @node is the right child of @parent.
289 * This function will adjust @parent's balance factor, then do a (single
290 * or double) rotation if necessary. The return value will be %true if
291 * the full AVL tree is now adequately balanced, or %false if the subtree
292 * rooted at @parent is now adequately balanced but has increased in
293 * height by 1, so the caller should continue up the tree.
295 * Note that if %false is returned, no rotation will have been done.
296 * Indeed, a single node insertion cannot require that more than one
297 * (single or double) rotation be done.
299 static AVL_INLINE bool
300 avl_handle_subtree_growth(struct avl_tree_node ** const root_ptr,
301 struct avl_tree_node * const node,
302 struct avl_tree_node * const parent,
305 int old_balance_factor, new_balance_factor;
307 old_balance_factor = avl_get_balance_factor(parent);
309 if (old_balance_factor == 0) {
310 avl_adjust_balance_factor(parent, sign);
311 /* @parent is still sufficiently balanced (-1 or +1
312 * balance factor), but must have increased in height.
313 * Continue up the tree. */
317 new_balance_factor = old_balance_factor + sign;
319 if (new_balance_factor == 0) {
320 avl_adjust_balance_factor(parent, sign);
321 /* @parent is now perfectly balanced (0 balance factor).
322 * It cannot have increased in height, so there is
323 * nothing more to do. */
327 /* @parent is too left-heavy (new_balance_factor == -2) or
328 * too right-heavy (new_balance_factor == +2). */
330 /* Test whether @node is left-heavy (-1 balance factor) or
331 * right-heavy (+1 balance factor).
332 * Note that it cannot be perfectly balanced (0 balance factor)
333 * because here we are under the invariant that @node has
334 * increased in height due to the insertion. */
335 if (sign * avl_get_balance_factor(node) > 0) {
337 /* @node (B below) is heavy in the same direction @parent
338 * (A below) is heavy.
340 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
341 * The comment, diagram, and equations below assume sign < 0.
342 * The other case is symmetric!
343 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
345 * Do a clockwise rotation rooted at @parent (A below):
355 * Before the rotation:
358 * Let x = height(C). Then:
362 * max(height(F), height(G)) = x.
364 * After the rotation:
365 * height(D) = max(height(F), height(G)) + 1
367 * height(A) = max(height(E), height(C)) + 1
368 * = max(x, x) + 1 = x + 1
372 avl_rotate(root_ptr, parent, -sign);
374 /* Equivalent to setting @parent's balance factor to 0. */
375 avl_adjust_balance_factor(parent, -sign); /* A */
377 /* Equivalent to setting @node's balance factor to 0. */
378 avl_adjust_balance_factor(node, -sign); /* B */
380 /* @node (B below) is heavy in the direction opposite
381 * from the direction @parent (A below) is heavy.
383 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
384 * The comment, diagram, and equations below assume sign < 0.
385 * The other case is symmetric!
386 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
388 * Do a counterblockwise rotation rooted at @node (B below),
389 * then a clockwise rotation rooted at @parent (A below):
393 * B C? => E C? => B A
395 * D? E B G? D? F?G? C?
399 * Before the rotation:
402 * Let x = height(C). Then:
406 * max(height(F), height(G)) = x
408 * After both rotations:
409 * height(A) = max(height(G), height(C)) + 1
411 * balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
412 * height(B) = max(height(D), height(F)) + 1
414 * balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
419 avl_do_double_rotate(root_ptr, node, parent, -sign);
422 /* Height after rotation is unchanged; nothing more to do. */
426 /* Rebalance the tree after insertion of the specified node. */
428 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
429 struct avl_tree_node *inserted)
431 struct avl_tree_node *node, *parent;
434 inserted->left = NULL;
435 inserted->right = NULL;
439 /* Adjust balance factor of new node's parent.
440 * No rotation will need to be done at this level. */
442 parent = avl_get_parent(node);
446 if (node == parent->left)
447 avl_adjust_balance_factor(parent, -1);
449 avl_adjust_balance_factor(parent, +1);
451 if (avl_get_balance_factor(parent) == 0)
452 /* @parent did not change in height. Nothing more to do. */
455 /* The subtree rooted at @parent increased in height by 1. */
458 /* Adjust balance factor of next ancestor. */
461 parent = avl_get_parent(node);
465 /* The subtree rooted at @node has increased in height by 1. */
466 if (node == parent->left)
467 done = avl_handle_subtree_growth(root_ptr, node,
470 done = avl_handle_subtree_growth(root_ptr, node,
476 * This function handles the shrinkage of a subtree due to a deletion.
479 * Location of the tree's root pointer.
482 * A node in the tree, exactly one of whose subtrees has decreased
483 * in height by 1 due to a deletion. (This includes the case where
484 * one of the child pointers has become NULL, since we can consider
485 * the "NULL" subtree to have a height of 0.)
488 * +1 if the left subtree of @parent has decreased in height by 1;
489 * -1 if the right subtree of @parent has decreased in height by 1.
492 * If the return value is not NULL, this will be set to %true if the
493 * left subtree of the returned node has decreased in height by 1,
494 * or %false if the right subtree of the returned node has decreased
497 * This function will adjust @parent's balance factor, then do a (single
498 * or double) rotation if necessary. The return value will be NULL if
499 * the full AVL tree is now adequately balanced, or a pointer to the
500 * parent of @parent if @parent is now adequately balanced but has
501 * decreased in height by 1. Also in the latter case, *left_deleted_ret
504 static AVL_INLINE struct avl_tree_node *
505 avl_handle_subtree_shrink(struct avl_tree_node ** const root_ptr,
506 struct avl_tree_node *parent,
508 bool * const left_deleted_ret)
510 struct avl_tree_node *node;
511 int old_balance_factor, new_balance_factor;
513 old_balance_factor = avl_get_balance_factor(parent);
515 if (old_balance_factor == 0) {
516 /* Prior to the deletion, the subtree rooted at
517 * @parent was perfectly balanced. It's now
518 * unbalanced by 1, but that's okay and its height
519 * hasn't changed. Nothing more to do. */
520 avl_adjust_balance_factor(parent, sign);
524 new_balance_factor = old_balance_factor + sign;
526 if (new_balance_factor == 0) {
527 /* The subtree rooted at @parent is now perfectly
528 * balanced, whereas before the deletion it was
529 * unbalanced by 1. Its height must have decreased
530 * by 1. No rotation is needed at this location,
531 * but continue up the tree. */
532 avl_adjust_balance_factor(parent, sign);
535 /* @parent is too left-heavy (new_balance_factor == -2) or
536 * too right-heavy (new_balance_factor == +2). */
538 node = avl_get_child(parent, sign);
540 /* The rotations below are similar to those done during
541 * insertion (see avl_handle_subtree_growth()), so full
542 * comments are not provided. The only new case is the
543 * one where @node has a balance factor of 0, and that is
546 if (sign * avl_get_balance_factor(node) >= 0) {
548 avl_rotate(root_ptr, parent, -sign);
550 if (avl_get_balance_factor(node) == 0) {
552 * @node (B below) is perfectly balanced.
554 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
555 * The comment, diagram, and equations
556 * below assume sign < 0. The other case
558 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
560 * Do a clockwise rotation rooted at
571 * Before the rotation:
574 * Let x = height(C). Then:
578 * max(height(F), height(G)) = x.
580 * After the rotation:
581 * height(D) = max(height(F), height(G)) + 1
583 * height(A) = max(height(E), height(C)) + 1
584 * = max(x + 1, x) + 1 = x + 2
589 /* A: -2 => -1 (sign < 0)
590 * or +2 => +1 (sign > 0)
591 * No change needed --- that's the same as
592 * old_balance_factor. */
594 /* B: 0 => +1 (sign < 0)
595 * or 0 => -1 (sign > 0) */
596 avl_adjust_balance_factor(node, -sign);
598 /* Height is unchanged; nothing more to do. */
601 avl_adjust_balance_factor(parent, -sign);
602 avl_adjust_balance_factor(node, -sign);
605 node = avl_do_double_rotate(root_ptr, node,
609 parent = avl_get_parent(node);
611 *left_deleted_ret = (node == parent->left);
615 /* Swaps node X, which must have 2 children, with its in-order successor, then
616 * unlinks node X. Returns the parent of X just before unlinking, without its
617 * balance factor having been updated to account for the unlink. */
618 static AVL_INLINE struct avl_tree_node *
619 avl_tree_swap_with_successor(struct avl_tree_node **root_ptr,
620 struct avl_tree_node *X,
621 bool *left_deleted_ret)
623 struct avl_tree_node *Y, *ret;
636 * [ X unlinked, Y returned ]
639 *left_deleted_ret = false;
641 struct avl_tree_node *Q;
653 * A ... => A ... => A ...
662 * [ X unlinked, Q returned ]
667 avl_set_parent(Q->left, Q);
669 avl_set_parent(X->right, Y);
671 *left_deleted_ret = true;
675 avl_set_parent(X->left, Y);
677 Y->parent_balance = X->parent_balance;
678 avl_replace_child(root_ptr, avl_get_parent(X), X, Y);
684 * Removes an item from the specified AVL tree.
687 * Location of the AVL tree's root pointer. Indirection is needed
688 * because the root node may change if the tree needed to be rebalanced
689 * because of the deletion or if @node was the root node.
692 * Pointer to the `struct avl_tree_node' embedded in the item to
693 * remove from the tree.
695 * Note: This function *only* removes the node and rebalances the tree.
696 * It does not free any memory, nor does it do the equivalent of
697 * avl_tree_node_set_unlinked().
700 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
702 struct avl_tree_node *parent;
703 bool left_deleted = false;
705 if (node->left && node->right) {
706 /* @node is fully internal, with two children. Swap it
707 * with its in-order successor (which must exist in the
708 * right subtree of @node and can have, at most, a right
709 * child), then unlink @node. */
710 parent = avl_tree_swap_with_successor(root_ptr, node,
712 /* @parent is now the parent of what was @node's in-order
713 * successor. It cannot be NULL, since @node itself was
714 * an ancestor of its in-order successor.
715 * @left_deleted has been set to %true if @node's
716 * in-order successor was the left child of @parent,
717 * otherwise %false. */
719 struct avl_tree_node *child;
721 /* @node is missing at least one child. Unlink it. Set
722 * @parent to @node's parent, and set @left_deleted to
723 * reflect which child of @parent @node was. Or, if
724 * @node was the root node, simply update the root node
726 child = node->left ? node->left : node->right;
727 parent = avl_get_parent(node);
729 if (node == parent->left) {
730 parent->left = child;
733 parent->right = child;
734 left_deleted = false;
737 avl_set_parent(child, parent);
740 avl_set_parent(child, parent);
746 /* Rebalance the tree. */
749 parent = avl_handle_subtree_shrink(root_ptr, parent,
752 parent = avl_handle_subtree_shrink(root_ptr, parent,