check in post-order test cases

2 files changed, 31 insertions, 6 deletions

diff --git a/src/main/TreeTraversal.java b/src/main/TreeTraversal.java
index fee494c..d69164e 100644
--- a/src/main/TreeTraversal.java
+++ b/src/main/TreeTraversal.java
@@ -125,8 +125,15 @@ public class TreeTraversal {
                         currentIndex = left[currentIndex];
                         keyIndexStack.push(currentIndex);
                     } else {
-                        result.add(key[keyIndexStack.pop()]);
-                        walkLeft = false;
+                        if (right[currentIndex] == -1) {
+                            result.add(key[currentIndex]);
+                            walkLeft = false;
+                            keyIndexStack.pop();
+                            currentIndex = keyIndexStack.peek();
+                        } else {
+                            currentIndex = right[currentIndex];
+                            keyIndexStack.push(currentIndex);
+                        }
                     }
                 } else {
                     if (right[currentIndex] != -1) {
@@ -134,10 +141,9 @@ public class TreeTraversal {
                         keyIndexStack.push(currentIndex);
                         walkLeft = true;
                     } else {
-                        if (!keyIndexStack.empty()) {
-                            currentIndex = keyIndexStack.pop();
-                            result.add(key[currentIndex]);
-                        }
+                        result.add(key[currentIndex]);
+                        keyIndexStack.pop();
+                        currentIndex = keyIndexStack.peek();
                     }
                 }
             }
diff --git a/src/test/TreeTraversalTest.java b/src/test/TreeTraversalTest.java
index 336b86c..1fd23ea 100644
--- a/src/test/TreeTraversalTest.java
+++ b/src/test/TreeTraversalTest.java
@@ -23,6 +23,13 @@ class TreeTraversalTest {
         assertEquals(1, tt.preOrder().get(2));
         assertEquals(3, tt.preOrder().get(3));
         assertEquals(5, tt.preOrder().get(4));
+
+        assertEquals(5, tt.postOrder().size());
+        assertEquals(1, tt.postOrder().get(0));
+        assertEquals(3, tt.postOrder().get(1));
+        assertEquals(2, tt.postOrder().get(2));
+        assertEquals(5, tt.postOrder().get(3));
+        assertEquals(4, tt.postOrder().get(4));
     }
 
     @Test
@@ -55,6 +62,18 @@ class TreeTraversalTest {
         assertEquals(20, tt.preOrder().get(7));
         assertEquals(60, tt.preOrder().get(8));
         assertEquals(10, tt.preOrder().get(9));
+
+        assertEquals(10, tt.postOrder().size());
+        assertEquals(50, tt.postOrder().get(0));
+        assertEquals(80, tt.postOrder().get(1));
+        assertEquals(90, tt.postOrder().get(2));
+        assertEquals(30, tt.postOrder().get(3));
+        assertEquals(40, tt.postOrder().get(4));
+        assertEquals(70, tt.postOrder().get(5));
+        assertEquals(10, tt.postOrder().get(6));
+        assertEquals(60, tt.postOrder().get(7));
+        assertEquals(20, tt.postOrder().get(8));
+        assertEquals(0, tt.postOrder().get(9));
     }
 
 }
\ No newline at end of file