BAEL-4864: Fix traverseInOrderWithoutRecursion method (eugenp#10588)

kwoyke · Krzysztof Woyke · web-flow · commit c6c90581a653 · 2021-03-24T22:09:02.000+01:00
Co-authored-by: Krzysztof Woyke &lt;krzysztof.woyke.sp@lhsystems.com&gt;
diff --git a/data-structures/src/main/java/com/baeldung/tree/BinaryTree.java b/data-structures/src/main/java/com/baeldung/tree/BinaryTree.java
@@ -148,48 +148,46 @@ public void traverseLevelOrder() {
         }
     }
 
-    
     public void traverseInOrderWithoutRecursion() {
-        Stack<Node> stack = new Stack<Node>();
+        Stack<Node> stack = new Stack<>();
         Node current = root;
-        stack.push(root);
-        while(! stack.isEmpty()) {
-            while(current.left != null) {
-                current = current.left;                
-                stack.push(current);                
-            }
-            current = stack.pop();
-            visit(current.value);
-            if(current.right != null) {
-                current = current.right;                
+
+        while (current != null || !stack.isEmpty()) {
+            while (current != null) {
                 stack.push(current);
+                current = current.left;
             }
+
+            Node top = stack.pop();
+            visit(top.value);
+            current = top.right;
         }
     }
-    
+
     public void traversePreOrderWithoutRecursion() {
-        Stack<Node> stack = new Stack<Node>();
+        Stack<Node> stack = new Stack<>();
         Node current = root;
         stack.push(root);
-        while(! stack.isEmpty()) {
+
+        while (current != null && !stack.isEmpty()) {
             current = stack.pop();
             visit(current.value);
-            
-            if(current.right != null)
+
+            if (current.right != null)
                 stack.push(current.right);
-                
-            if(current.left != null)
+
+            if (current.left != null)
                 stack.push(current.left);
-        }        
+        }
     }
     
     public void traversePostOrderWithoutRecursion() {
-        Stack<Node> stack = new Stack<Node>();
+        Stack<Node> stack = new Stack<>();
         Node prev = root;
         Node current = root;
         stack.push(root);
 
-        while (!stack.isEmpty()) {
+        while (current != null && !stack.isEmpty()) {
             current = stack.peek();
             boolean hasChild = (current.left != null || current.right != null);
             boolean isPrevLastChild = (prev == current.right || (prev == current.left && current.right == null));
diff --git a/data-structures/src/test/java/com/baeldung/tree/BinaryTreeUnitTest.java b/data-structures/src/test/java/com/baeldung/tree/BinaryTreeUnitTest.java
@@ -13,7 +13,7 @@ public void givenABinaryTree_WhenAddingElements_ThenTreeNotEmpty() {
 
         BinaryTree bt = createBinaryTree();
 
-        assertTrue(!bt.isEmpty());
+        assertFalse(bt.isEmpty());
     }
 
     @Test
@@ -72,6 +72,7 @@ public void givenABinaryTree_WhenDeletingNonExistingElement_ThenTreeDoesNotDelet
 
     @Test
     public void it_deletes_the_root() {
+
         int value = 12;
         BinaryTree bt = new BinaryTree();
         bt.add(value);
@@ -91,6 +92,14 @@ public void givenABinaryTree_WhenTraversingInOrder_ThenPrintValues() {
         bt.traverseInOrderWithoutRecursion();
     }
 
+    @Test
+    public void givenAnEmptyBinaryTree_WhenTraversingInOrderWithoutRecursion_ThenNoException() {
+
+        BinaryTree empty = new BinaryTree();
+
+        empty.traverseInOrderWithoutRecursion();
+    }
+
     @Test
     public void givenABinaryTree_WhenTraversingPreOrder_ThenPrintValues() {
 
@@ -101,6 +110,14 @@ public void givenABinaryTree_WhenTraversingPreOrder_ThenPrintValues() {
         bt.traversePreOrderWithoutRecursion();
     }
 
+    @Test
+    public void givenAnEmptyBinaryTree_WhenTraversingPreOrderWithoutRecursion_ThenNoException() {
+
+        BinaryTree empty = new BinaryTree();
+
+        empty.traversePreOrderWithoutRecursion();
+    }
+
     @Test
     public void givenABinaryTree_WhenTraversingPostOrder_ThenPrintValues() {
 
@@ -111,6 +128,14 @@ public void givenABinaryTree_WhenTraversingPostOrder_ThenPrintValues() {
         bt.traversePostOrderWithoutRecursion();
     }
 
+    @Test
+    public void givenAnEmptyBinaryTree_WhenTraversingPostOrderWithoutRecursion_ThenNoException() {
+
+        BinaryTree empty = new BinaryTree();
+
+        empty.traversePostOrderWithoutRecursion();
+    }
+
     @Test
     public void givenABinaryTree_WhenTraversingLevelOrder_ThenPrintValues() {
 
