tree traversal done (in/pre/post) orders

Not too bad since I worked it out in Java first

2 files changed, 163 insertions, 0 deletions

diff --git a/sources/tree_traversal.py b/sources/tree_traversal.py
new file mode 100644
index 0000000..b75b0f7
--- /dev/null
+++ b/sources/tree_traversal.py
@@ -0,0 +1,126 @@
+# python3
+import queue
+import sys
+
+
+def in_order(key, left, right):
+    keyIndexStack = []
+    result = []
+    walkLeft = True
+
+    currentIndex = 0
+    keyIndexStack.append(currentIndex)
+    while len(keyIndexStack) != 0 or right[currentIndex] != -1:
+        if walkLeft:
+            if left[currentIndex] != -1:
+                currentIndex = left[currentIndex]
+                keyIndexStack.append(currentIndex)
+            else:
+                result.append(key[keyIndexStack.pop()])
+                walkLeft = False
+        else:
+            if right[currentIndex] != -1:
+                currentIndex = right[currentIndex]
+                keyIndexStack.append(currentIndex)
+                walkLeft = True
+            else:
+                if len(keyIndexStack) != 0:
+                    currentIndex = keyIndexStack.pop()
+                    result.append(key[currentIndex])
+    return result
+
+
+def pre_order(key, left, right):
+    keyIndexStack = []
+    keyIndexQueue = queue.Queue()
+    result = []
+
+    walkLeft = True
+    currentIndex = 0
+    keyIndexStack.append(currentIndex)
+    keyIndexQueue.put(currentIndex)
+
+    while len(keyIndexStack) != 0 or right[currentIndex] != -1:
+        if walkLeft:
+            if left[currentIndex] != -1:
+                currentIndex = left[currentIndex]
+                keyIndexQueue.put(currentIndex)
+                keyIndexStack.append(currentIndex)
+            else:
+                keyIndexStack.pop()
+                walkLeft = False
+                result.append(key[keyIndexQueue.get()])
+        else:
+            if right[currentIndex] != -1:
+                currentIndex = right[currentIndex]
+                keyIndexQueue.put(currentIndex)
+                keyIndexStack.append(currentIndex)
+                walkLeft = True
+            else:
+                if len(keyIndexStack) != 0:
+                    currentIndex = keyIndexStack.pop()
+                    result.append(key[keyIndexQueue.get()])
+    return result
+
+
+def post_order(key, left, right):
+    keyIndexStack = []
+    peekedSet = set([])
+    result = []
+
+    walkLeft = True
+    currentIndex = 0
+    keyIndexStack.append(currentIndex)
+    while len(keyIndexStack) != 0:
+        if currentIndex in peekedSet:
+            result.append(key[currentIndex])
+            peekedSet.remove(currentIndex)
+            keyIndexStack.pop()
+            if len(keyIndexStack) == 0:
+                break
+            currentIndex = keyIndexStack[-1]
+            continue
+        if walkLeft:
+            if left[currentIndex] != -1:
+                currentIndex = left[currentIndex]
+                keyIndexStack.append(currentIndex)
+            else:
+                if right[currentIndex] == -1:
+                    result.append(key[currentIndex])
+                    walkLeft = False
+                    keyIndexStack.pop()
+                    currentIndex = keyIndexStack[-1]
+                else:
+                    peekedSet.add(currentIndex)
+                    currentIndex = right[currentIndex]
+                    keyIndexStack.append(currentIndex)
+        else:
+            if right[currentIndex] != -1:
+                peekedSet.add(currentIndex)
+                currentIndex = right[currentIndex]
+                keyIndexStack.append(currentIndex)
+                walkLeft = True
+            else:
+                result.append(key[currentIndex])
+                i = keyIndexStack.pop();
+                if i != currentIndex:
+                    result.append(key[i])
+                if len(keyIndexStack) == 0:
+                    break
+                currentIndex = keyIndexStack[-1]
+    return result
+
+
+if __name__ == '__main__':
+    n = int(sys.stdin.readline())
+    key = [0 for i in range(n)]
+    left = [0 for i in range(n)]
+    right = [0 for i in range(n)]
+    for i in range(n):
+        [a, b, c] = map(int, sys.stdin.readline().split())
+        key[i] = a
+        left[i] = b
+        right[i] = c
+    print(" ".join(str(x) for x in in_order(key, left, right)))
+    print(" ".join(str(x) for x in pre_order(key, left, right)))
+    print(" ".join(str(x) for x in post_order(key, left, right)))
diff --git a/tests/tree_traversalTest.py b/tests/tree_traversalTest.py
new file mode 100644
index 0000000..7fe79e3
--- /dev/null
+++ b/tests/tree_traversalTest.py
@@ -0,0 +1,37 @@
+import unittest
+
+from sources.tree_traversal import in_order, pre_order, post_order
+
+class MyTestCase(unittest.TestCase):
+    def test_something(self):
+        key = [4, 2, 5, 1, 3]
+        left = [1, 3, -1, -1, -1]
+        right = [2, 4, -1, -1, -1]
+
+        result = in_order(key, left, right)
+        self.assertEqual(len(result), 5)
+        self.assertEqual(result[0], 1)
+        self.assertEqual(result[1], 2)
+        self.assertEqual(result[2], 3)
+        self.assertEqual(result[3], 4)
+        self.assertEqual(result[4], 5)
+
+        result = pre_order(key, left, right)
+        self.assertEqual(len(result), 5)
+        self.assertEqual(result[0], 4)
+        self.assertEqual(result[1], 2)
+        self.assertEqual(result[2], 1)
+        self.assertEqual(result[3], 3)
+        self.assertEqual(result[4], 5)
+
+        result = post_order(key, left, right)
+        self.assertEqual(len(result), 5)
+        self.assertEqual(result[0], 1)
+        self.assertEqual(result[1], 3)
+        self.assertEqual(result[2], 2)
+        self.assertEqual(result[3], 5)
+        self.assertEqual(result[4], 4)
+
+
+if __name__ == '__main__':
+    unittest.main()