diff --git a/avl_tree.c b/avl_tree.c index 4fdff6f..e5b911d 100644 --- a/avl_tree.c +++ b/avl_tree.c @@ -787,3 +787,35 @@ avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node) -1, &left_deleted); } while (parent); } + +/* Starts a preorder traversal of the tree. */ +struct avl_tree_node * +avl_tree_first_in_preorder(const struct avl_tree_node *root) +{ + return (struct avl_tree_node *)root; +} + +/* Continues a preorder traversal of the tree. @prev_parent must be its saved + * parent node. Returns NULL if there are no more nodes (i.e. @prev was the + * root of the tree). */ +struct avl_tree_node * +avl_tree_next_in_preorder(const struct avl_tree_node *prev, + const struct avl_tree_node *prev_parent) +{ + const struct avl_tree_node *next = NULL; + + if (prev && (prev->left || prev->right)) + next = prev->left ? prev->left : prev->right; + else + for (; + prev_parent; + prev = prev_parent, prev_parent = avl_get_parent(prev_parent)) + if (prev_parent->right && prev_parent->right != prev) + { + next = prev_parent->right; + break; + } + + return (struct avl_tree_node *)next; +} + diff --git a/avl_tree.h b/avl_tree.h index 0dc6b4e..8b84c43 100644 --- a/avl_tree.h +++ b/avl_tree.h @@ -285,6 +285,13 @@ avl_tree_next_in_order(const struct avl_tree_node *node); extern struct avl_tree_node * avl_tree_prev_in_order(const struct avl_tree_node *node); +extern struct avl_tree_node * +avl_tree_first_in_preorder(const struct avl_tree_node *root); + +extern struct avl_tree_node * +avl_tree_next_in_preorder(const struct avl_tree_node *prev, + const struct avl_tree_node *prev_parent); + extern struct avl_tree_node * avl_tree_first_in_postorder(const struct avl_tree_node *root); @@ -342,18 +349,33 @@ avl_tree_next_in_postorder(const struct avl_tree_node *prev, struct_member), 1); \ _cur = avl_tree_prev_in_order(_cur)) + +#define avl_tree_for_each(child_struct, root, \ + struct_name, struct_member, order) \ + for (struct avl_tree_node *_cur = \ + avl_tree_first_in_##order(root), *_parent; \ + _cur && ((child_struct) = \ + avl_tree_entry(_cur, struct_name, \ + struct_member), 1) \ + && (_parent = avl_get_parent(_cur), 1); \ + _cur = avl_tree_next_in_##order(_cur, _parent)) + /* * Like avl_tree_for_each_in_order(), but iterates through the nodes in * postorder, so the current node may be deleted or freed. */ #define avl_tree_for_each_in_postorder(child_struct, root, \ + struct_name, struct_member) \ + avl_tree_for_each(child_struct, root, \ + struct_name, struct_member, postorder) + +/* + * Like avl_tree_for_each_in_order(), but iterates through the nodes in + * preorder. + */ +#define avl_tree_for_each_in_preorder(child_struct, root, \ struct_name, struct_member) \ - for (struct avl_tree_node *_cur = \ - avl_tree_first_in_postorder(root), *_parent; \ - _cur && ((child_struct) = \ - avl_tree_entry(_cur, struct_name, \ - struct_member), 1) \ - && (_parent = avl_get_parent(_cur), 1); \ - _cur = avl_tree_next_in_postorder(_cur, _parent)) + avl_tree_for_each(child_struct, root, \ + struct_name, struct_member, preorder) #endif /* _AVL_TREE_H_ */ diff --git a/test.c b/test.c index 55c6bdb..82ce44e 100644 --- a/test.c +++ b/test.c @@ -192,6 +192,20 @@ verify(const int *data, int count) assert(!cur->right || TEST_NODE(cur->right)->reached); } assert(x == count); + + /* Check preorder traversal. */ + for (cur = avl_tree_first_in_preorder(root), x = 0; + cur; + cur = avl_tree_next_in_preorder(cur, avl_get_parent(cur)), x++) + { + assert(TEST_NODE(cur)->reached); + TEST_NODE(cur)->reached = 0; + assert(!avl_get_parent(cur) || + !(TEST_NODE(avl_get_parent(cur))->reached)); + assert(!cur->left || TEST_NODE(cur->left)->reached); + assert(!cur->right || TEST_NODE(cur->right)->reached); + } + assert(x == count); } #endif