1 files changed, 25 insertions, 9 deletions

diff --git a/src/main/TreeTraversal.java b/src/main/TreeTraversal.java
index d69164e..d7ca19d 100644
--- a/src/main/TreeTraversal.java
+++ b/src/main/TreeTraversal.java
@@ -23,7 +23,7 @@ public class TreeTraversal {
         }
     }
 
-    public static class TreeOrders {
+    static class TreeOrders {
         int n;
         int[] key, left, right;
 
@@ -42,7 +42,7 @@ public class TreeTraversal {
 
         List<Integer> inOrder() {
             Stack<Integer> keyIndexStack = new Stack<>();
-            ArrayList<Integer> result = new ArrayList<Integer>();
+            ArrayList<Integer> result = new ArrayList<>();
             // Finish the implementation
             // You may need to add a new recursive method to do that
             boolean walkLeft = true;
@@ -76,7 +76,7 @@ public class TreeTraversal {
         List<Integer> preOrder() {
             Stack<Integer> keyIndexStack = new Stack<>();
             Queue<Integer> keyIndexQueue = new LinkedList<>();
-            ArrayList<Integer> result = new ArrayList<Integer>();
+            ArrayList<Integer> result = new ArrayList<>();
             // Finish the implementation
             // You may need to add a new recursive method to do that
             boolean walkLeft = true;
@@ -113,13 +113,24 @@ public class TreeTraversal {
 
         List<Integer> postOrder() {
             Stack<Integer> keyIndexStack = new Stack<>();
-            ArrayList<Integer> result = new ArrayList<Integer>();
+            HashSet<Integer> peekedSet = new HashSet<>();
+            ArrayList<Integer> result = new ArrayList<>();
+//            ArrayList<Integer> bothSidesHaveBeenChecked = new ArrayList<>();
             // Finish the implementation
             // You may need to add a new recursive method to do that
             boolean walkLeft = true;
             int currentIndex = 0;
             keyIndexStack.push(currentIndex);
-            while (!keyIndexStack.empty() || right[currentIndex] != -1) {
+//            while (!keyIndexStack.empty() || right[currentIndex] != -1) {
+            while (!keyIndexStack.empty()) {
+                if (peekedSet.contains(currentIndex)) {
+                    result.add(key[currentIndex]);
+                    peekedSet.remove(currentIndex);
+                    keyIndexStack.pop();
+                    if (keyIndexStack.empty()) break;
+                    currentIndex = keyIndexStack.peek();
+                    continue;
+                }
                 if (walkLeft) {
                     if (left[currentIndex] != -1) {
                         currentIndex = left[currentIndex];
@@ -131,18 +142,23 @@ public class TreeTraversal {
                             keyIndexStack.pop();
                             currentIndex = keyIndexStack.peek();
                         } else {
+                            peekedSet.add(currentIndex);
                             currentIndex = right[currentIndex];
                             keyIndexStack.push(currentIndex);
                         }
                     }
                 } else {
                     if (right[currentIndex] != -1) {
+                        peekedSet.add(currentIndex);
                         currentIndex = right[currentIndex];
                         keyIndexStack.push(currentIndex);
                         walkLeft = true;
                     } else {
                         result.add(key[currentIndex]);
-                        keyIndexStack.pop();
+                        int i = keyIndexStack.pop();
+                        if (i != currentIndex)
+                            result.add(key[i]);
+                        if (keyIndexStack.empty()) break;
                         currentIndex = keyIndexStack.peek();
                     }
                 }
@@ -156,20 +172,20 @@ public class TreeTraversal {
             public void run() {
                 try {
                     new TreeTraversal().run();
-                } catch (IOException e) {
+                } catch (IOException ignored) {
                 }
             }
         }, "1", 1 << 26).start();
     }
 
-    public void print(List<Integer> x) {
+    private void print(List<Integer> x) {
         for (Integer a : x) {
             System.out.print(a + " ");
         }
         System.out.println();
     }
 
-    public void run() throws IOException {
+    private void run() throws IOException {
         TreeOrders tree = new TreeOrders();
         tree.read();
         print(tree.inOrder());