diff options
| -rw-r--r-- | sources/tree_traversal.py | 126 | ||||
| -rw-r--r-- | tests/tree_traversalTest.py | 37 |
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() |