1 files changed, 33 insertions, 1 deletions

diff --git a/src/02/01/binary_tree_test.c b/src/02/01/binary_tree_test.c
index f1856b6..db0a705 100644
--- a/src/02/01/binary_tree_test.c
+++ b/src/02/01/binary_tree_test.c
@@ -2,21 +2,53 @@
 #include <cgreen/cgreen.h>
 #include <string.h>
 
-int visited[256];
+int visited[32];
+Node *nodes[32];
 int visited_count;
 
 Describe(BinaryTree);
 BeforeEach(BinaryTree) {
   memset(visited, 0, sizeof(visited));
+  memset(nodes, 0, sizeof(nodes));
   visited_count = 0;
 }
 AfterEach(BinaryTree) {}
 
 void visitor(Node *node) {
   visited[visited_count] = node->data;
+  nodes[visited_count] = node;
   visited_count++;
 }
 
+Node *preorder_next(Node *self, Node *x) {
+  traverse(self, visitor, PREORDER);
+
+  for (int i = 0; i < visited_count; i++)
+    if (nodes[i] == x)
+      return nodes[i + 1];
+  return NULL;
+}
+
+Ensure(BinaryTree, when_finding_the_next_node_in_a_preorder_traversal) {
+  Node *a = initialize(100);
+  Node *b = initialize(200);
+  Node *c = initialize(300);
+  Node *d = initialize(400);
+  Node *e = initialize(500);
+
+  a->left = b;
+  a->right = c;
+  b->left = d;
+  b->right = e;
+
+  assert_that(preorder_next(a, a), is_equal_to(b));
+  assert_that(preorder_next(a, b), is_equal_to(d));
+  assert_that(preorder_next(a, d), is_equal_to(e));
+  assert_that(preorder_next(a, e), is_equal_to(c));
+
+  destroy(a);
+}
+
 Ensure(BinaryTree, when_traversing_in_preorder_when_the_tree_is_empty) {
   traverse(NULL, visitor, PREORDER);