TheAlgorithms-Python

1
"""
2
This is pure Python implementation of tree traversal algorithms
3
"""
4

5
from __future__ import annotations
6

7
import queue
8

9

10
class TreeNode:
11
    def __init__(self, data):
12
        self.data = data
13
        self.right = None
14
        self.left = None
15

16

17
def build_tree() -> TreeNode:
18
    print("\n********Press N to stop entering at any point of time********\n")
19
    check = input("Enter the value of the root node: ").strip().lower()
20
    q: queue.Queue = queue.Queue()
21
    tree_node = TreeNode(int(check))
22
    q.put(tree_node)
23
    while not q.empty():
24
        node_found = q.get()
25
        msg = f"Enter the left node of {node_found.data}: "
26
        check = input(msg).strip().lower() or "n"
27
        if check == "n":
28
            return tree_node
29
        left_node = TreeNode(int(check))
30
        node_found.left = left_node
31
        q.put(left_node)
32
        msg = f"Enter the right node of {node_found.data}: "
33
        check = input(msg).strip().lower() or "n"
34
        if check == "n":
35
            return tree_node
36
        right_node = TreeNode(int(check))
37
        node_found.right = right_node
38
        q.put(right_node)
39
    raise
40

41

42
def pre_order(node: TreeNode) -> None:
43
    """
44
    >>> root = TreeNode(1)
45
    >>> tree_node2 = TreeNode(2)
46
    >>> tree_node3 = TreeNode(3)
47
    >>> tree_node4 = TreeNode(4)
48
    >>> tree_node5 = TreeNode(5)
49
    >>> tree_node6 = TreeNode(6)
50
    >>> tree_node7 = TreeNode(7)
51
    >>> root.left, root.right = tree_node2, tree_node3
52
    >>> tree_node2.left, tree_node2.right = tree_node4 , tree_node5
53
    >>> tree_node3.left, tree_node3.right = tree_node6 , tree_node7
54
    >>> pre_order(root)
55
    1,2,4,5,3,6,7,
56
    """
57
    if not isinstance(node, TreeNode) or not node:
58
        return
59
    print(node.data, end=",")
60
    pre_order(node.left)
61
    pre_order(node.right)
62

63

64
def in_order(node: TreeNode) -> None:
65
    """
66
    >>> root = TreeNode(1)
67
    >>> tree_node2 = TreeNode(2)
68
    >>> tree_node3 = TreeNode(3)
69
    >>> tree_node4 = TreeNode(4)
70
    >>> tree_node5 = TreeNode(5)
71
    >>> tree_node6 = TreeNode(6)
72
    >>> tree_node7 = TreeNode(7)
73
    >>> root.left, root.right = tree_node2, tree_node3
74
    >>> tree_node2.left, tree_node2.right = tree_node4 , tree_node5
75
    >>> tree_node3.left, tree_node3.right = tree_node6 , tree_node7
76
    >>> in_order(root)
77
    4,2,5,1,6,3,7,
78
    """
79
    if not isinstance(node, TreeNode) or not node:
80
        return
81
    in_order(node.left)
82
    print(node.data, end=",")
83
    in_order(node.right)
84

85

86
def post_order(node: TreeNode) -> None:
87
    """
88
    >>> root = TreeNode(1)
89
    >>> tree_node2 = TreeNode(2)
90
    >>> tree_node3 = TreeNode(3)
91
    >>> tree_node4 = TreeNode(4)
92
    >>> tree_node5 = TreeNode(5)
93
    >>> tree_node6 = TreeNode(6)
94
    >>> tree_node7 = TreeNode(7)
95
    >>> root.left, root.right = tree_node2, tree_node3
96
    >>> tree_node2.left, tree_node2.right = tree_node4 , tree_node5
97
    >>> tree_node3.left, tree_node3.right = tree_node6 , tree_node7
98
    >>> post_order(root)
99
    4,5,2,6,7,3,1,
100
    """
101
    if not isinstance(node, TreeNode) or not node:
102
        return
103
    post_order(node.left)
104
    post_order(node.right)
105
    print(node.data, end=",")
106

107

108
def level_order(node: TreeNode) -> None:
109
    """
110
    >>> root = TreeNode(1)
111
    >>> tree_node2 = TreeNode(2)
112
    >>> tree_node3 = TreeNode(3)
113
    >>> tree_node4 = TreeNode(4)
114
    >>> tree_node5 = TreeNode(5)
115
    >>> tree_node6 = TreeNode(6)
116
    >>> tree_node7 = TreeNode(7)
117
    >>> root.left, root.right = tree_node2, tree_node3
118
    >>> tree_node2.left, tree_node2.right = tree_node4 , tree_node5
119
    >>> tree_node3.left, tree_node3.right = tree_node6 , tree_node7
120
    >>> level_order(root)
121
    1,2,3,4,5,6,7,
122
    """
123
    if not isinstance(node, TreeNode) or not node:
124
        return
125
    q: queue.Queue = queue.Queue()
126
    q.put(node)
127
    while not q.empty():
128
        node_dequeued = q.get()
129
        print(node_dequeued.data, end=",")
130
        if node_dequeued.left:
131
            q.put(node_dequeued.left)
132
        if node_dequeued.right:
133
            q.put(node_dequeued.right)
134

135

136
def level_order_actual(node: TreeNode) -> None:
137
    """
138
    >>> root = TreeNode(1)
139
    >>> tree_node2 = TreeNode(2)
140
    >>> tree_node3 = TreeNode(3)
141
    >>> tree_node4 = TreeNode(4)
142
    >>> tree_node5 = TreeNode(5)
143
    >>> tree_node6 = TreeNode(6)
144
    >>> tree_node7 = TreeNode(7)
145
    >>> root.left, root.right = tree_node2, tree_node3
146
    >>> tree_node2.left, tree_node2.right = tree_node4 , tree_node5
147
    >>> tree_node3.left, tree_node3.right = tree_node6 , tree_node7
148
    >>> level_order_actual(root)
149
    1,
150
    2,3,
151
    4,5,6,7,
152
    """
153
    if not isinstance(node, TreeNode) or not node:
154
        return
155
    q: queue.Queue = queue.Queue()
156
    q.put(node)
157
    while not q.empty():
158
        list_ = []
159
        while not q.empty():
160
            node_dequeued = q.get()
161
            print(node_dequeued.data, end=",")
162
            if node_dequeued.left:
163
                list_.append(node_dequeued.left)
164
            if node_dequeued.right:
165
                list_.append(node_dequeued.right)
166
        print()
167
        for node in list_:
168
            q.put(node)
169

170

171
# iteration version
172
def pre_order_iter(node: TreeNode) -> None:
173
    """
174
    >>> root = TreeNode(1)
175
    >>> tree_node2 = TreeNode(2)
176
    >>> tree_node3 = TreeNode(3)
177
    >>> tree_node4 = TreeNode(4)
178
    >>> tree_node5 = TreeNode(5)
179
    >>> tree_node6 = TreeNode(6)
180
    >>> tree_node7 = TreeNode(7)
181
    >>> root.left, root.right = tree_node2, tree_node3
182
    >>> tree_node2.left, tree_node2.right = tree_node4 , tree_node5
183
    >>> tree_node3.left, tree_node3.right = tree_node6 , tree_node7
184
    >>> pre_order_iter(root)
185
    1,2,4,5,3,6,7,
186
    """
187
    if not isinstance(node, TreeNode) or not node:
188
        return
189
    stack: list[TreeNode] = []
190
    n = node
191
    while n or stack:
192
        while n:  # start from root node, find its left child
193
            print(n.data, end=",")
194
            stack.append(n)
195
            n = n.left
196
        # end of while means current node doesn't have left child
197
        n = stack.pop()
198
        # start to traverse its right child
199
        n = n.right
200

201

202
def in_order_iter(node: TreeNode) -> None:
203
    """
204
    >>> root = TreeNode(1)
205
    >>> tree_node2 = TreeNode(2)
206
    >>> tree_node3 = TreeNode(3)
207
    >>> tree_node4 = TreeNode(4)
208
    >>> tree_node5 = TreeNode(5)
209
    >>> tree_node6 = TreeNode(6)
210
    >>> tree_node7 = TreeNode(7)
211
    >>> root.left, root.right = tree_node2, tree_node3
212
    >>> tree_node2.left, tree_node2.right = tree_node4 , tree_node5
213
    >>> tree_node3.left, tree_node3.right = tree_node6 , tree_node7
214
    >>> in_order_iter(root)
215
    4,2,5,1,6,3,7,
216
    """
217
    if not isinstance(node, TreeNode) or not node:
218
        return
219
    stack: list[TreeNode] = []
220
    n = node
221
    while n or stack:
222
        while n:
223
            stack.append(n)
224
            n = n.left
225
        n = stack.pop()
226
        print(n.data, end=",")
227
        n = n.right
228

229

230
def post_order_iter(node: TreeNode) -> None:
231
    """
232
    >>> root = TreeNode(1)
233
    >>> tree_node2 = TreeNode(2)
234
    >>> tree_node3 = TreeNode(3)
235
    >>> tree_node4 = TreeNode(4)
236
    >>> tree_node5 = TreeNode(5)
237
    >>> tree_node6 = TreeNode(6)
238
    >>> tree_node7 = TreeNode(7)
239
    >>> root.left, root.right = tree_node2, tree_node3
240
    >>> tree_node2.left, tree_node2.right = tree_node4 , tree_node5
241
    >>> tree_node3.left, tree_node3.right = tree_node6 , tree_node7
242
    >>> post_order_iter(root)
243
    4,5,2,6,7,3,1,
244
    """
245
    if not isinstance(node, TreeNode) or not node:
246
        return
247
    stack1, stack2 = [], []
248
    n = node
249
    stack1.append(n)
250
    while stack1:  # to find the reversed order of post order, store it in stack2
251
        n = stack1.pop()
252
        if n.left:
253
            stack1.append(n.left)
254
        if n.right:
255
            stack1.append(n.right)
256
        stack2.append(n)
257
    while stack2:  # pop up from stack2 will be the post order
258
        print(stack2.pop().data, end=",")
259

260

261
def prompt(s: str = "", width=50, char="*") -> str:
262
    if not s:
263
        return "\n" + width * char
264
    left, extra = divmod(width - len(s) - 2, 2)
265
    return f"{left * char} {s} {(left + extra) * char}"
266

267

268
if __name__ == "__main__":
269
    import doctest
270

271
    doctest.testmod()
272
    print(prompt("Binary Tree Traversals"))
273

274
    node: TreeNode = build_tree()
275
    print(prompt("Pre Order Traversal"))
276
    pre_order(node)
277
    print(prompt() + "\n")
278

279
    print(prompt("In Order Traversal"))
280
    in_order(node)
281
    print(prompt() + "\n")
282

283
    print(prompt("Post Order Traversal"))
284
    post_order(node)
285
    print(prompt() + "\n")
286

287
    print(prompt("Level Order Traversal"))
288
    level_order(node)
289
    print(prompt() + "\n")
290

291
    print(prompt("Actual Level Order Traversal"))
292
    level_order_actual(node)
293
    print("*" * 50 + "\n")
294

295
    print(prompt("Pre Order Traversal - Iteration Version"))
296
    pre_order_iter(node)
297
    print(prompt() + "\n")
298

299
    print(prompt("In Order Traversal - Iteration Version"))
300
    in_order_iter(node)
301
    print(prompt() + "\n")
302

303
    print(prompt("Post Order Traversal - Iteration Version"))
304
    post_order_iter(node)
305
    print(prompt())
306