1
1
from __future__ import annotations
2
2
3
3
from collections import deque
4
- from collections .abc import Generator , Sequence
4
+ from collections .abc import Generator
5
5
from dataclasses import dataclass
6
- from typing import Any
7
6
8
7
9
8
# https://en.wikipedia.org/wiki/Tree_traversal
@@ -94,96 +93,86 @@ def height(root: Node | None) -> int:
94
93
return (max (height (root .left ), height (root .right )) + 1 ) if root else 0
95
94
96
95
97
- def level_order (root : Node | None ) -> Sequence [ Node | None ]:
96
+ def level_order (root : Node | None ) -> Generator [ int , None , None ]:
98
97
"""
99
98
Returns a list of nodes value from a whole binary tree in Level Order Traverse.
100
99
Level Order traverse: Visit nodes of the tree level-by-level.
101
100
"""
102
- output : list [Any ] = []
103
101
104
102
if root is None :
105
- return output
103
+ return
106
104
107
105
process_queue = deque ([root ])
108
106
109
107
while process_queue :
110
108
node = process_queue .popleft ()
111
- output . append ( node .data )
109
+ yield node .data
112
110
113
111
if node .left :
114
112
process_queue .append (node .left )
115
113
if node .right :
116
114
process_queue .append (node .right )
117
- return output
118
115
119
116
120
117
def get_nodes_from_left_to_right (
121
118
root : Node | None , level : int
122
- ) -> Sequence [ Node | None ]:
119
+ ) -> Generator [ int , None , None ]:
123
120
"""
124
121
Returns a list of nodes value from a particular level:
125
122
Left to right direction of the binary tree.
126
123
"""
127
- output : list [Any ] = []
128
124
129
- def populate_output (root : Node | None , level : int ) -> None :
125
+ def populate_output (root : Node | None , level : int ) -> Generator [ int , None , None ] :
130
126
if not root :
131
127
return
132
128
if level == 1 :
133
- output . append ( root .data )
129
+ yield root .data
134
130
elif level > 1 :
135
- populate_output (root .left , level - 1 )
136
- populate_output (root .right , level - 1 )
131
+ yield from populate_output (root .left , level - 1 )
132
+ yield from populate_output (root .right , level - 1 )
137
133
138
- populate_output (root , level )
139
- return output
134
+ yield from populate_output (root , level )
140
135
141
136
142
137
def get_nodes_from_right_to_left (
143
138
root : Node | None , level : int
144
- ) -> Sequence [ Node | None ]:
139
+ ) -> Generator [ int , None , None ]:
145
140
"""
146
141
Returns a list of nodes value from a particular level:
147
142
Right to left direction of the binary tree.
148
143
"""
149
- output : list [Any ] = []
150
144
151
- def populate_output (root : Node | None , level : int ) -> None :
145
+ def populate_output (root : Node | None , level : int ) -> Generator [ int , None , None ] :
152
146
if root is None :
153
147
return
154
148
if level == 1 :
155
- output . append ( root .data )
149
+ yield root .data
156
150
elif level > 1 :
157
- populate_output (root .right , level - 1 )
158
- populate_output (root .left , level - 1 )
151
+ yield from populate_output (root .right , level - 1 )
152
+ yield from populate_output (root .left , level - 1 )
159
153
160
- populate_output (root , level )
161
- return output
154
+ yield from populate_output (root , level )
162
155
163
156
164
- def zigzag (root : Node | None ) -> Sequence [ Node | None ] | list [ Any ]:
157
+ def zigzag (root : Node | None ) -> Generator [ int , None , None ]:
165
158
"""
166
159
ZigZag traverse:
167
160
Returns a list of nodes value from left to right and right to left, alternatively.
168
161
"""
169
162
if root is None :
170
- return []
171
-
172
- output : list [Sequence [Node | None ]] = []
163
+ return
173
164
174
165
flag = 0
175
166
height_tree = height (root )
176
167
177
168
for h in range (1 , height_tree + 1 ):
178
169
if not flag :
179
- output . append ( get_nodes_from_left_to_right (root , h ) )
170
+ yield from get_nodes_from_left_to_right (root , h )
180
171
flag = 1
181
172
else :
182
- output . append ( get_nodes_from_right_to_left (root , h ) )
173
+ yield from get_nodes_from_right_to_left (root , h )
183
174
flag = 0
184
175
185
- return output
186
-
187
176
188
177
def main () -> None : # Main function for testing.
189
178
# Create binary tree.
@@ -198,15 +187,15 @@ def main() -> None: # Main function for testing.
198
187
print (f"Height of Tree: { height (root )} " , "\n " )
199
188
200
189
print ("Complete Level Order Traversal: " )
201
- print (level_order (root ), " \n " )
190
+ print (f" { list ( level_order (root )) } \n " )
202
191
203
192
print ("Level-wise order Traversal: " )
204
193
205
194
for level in range (1 , height (root ) + 1 ):
206
- print (f"Level { level } :" , get_nodes_from_left_to_right (root , level = level ))
195
+ print (f"Level { level } :" , list ( get_nodes_from_left_to_right (root , level = level ) ))
207
196
208
197
print ("\n ZigZag order Traversal: " )
209
- print (zigzag (root ))
198
+ print (f" { list ( zigzag (root )) } " )
210
199
211
200
212
201
if __name__ == "__main__" :
0 commit comments