summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--sources/tree_traversal.py126
-rw-r--r--tests/tree_traversalTest.py37
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()