11from __future__ import annotations
22
33from collections import deque
4- from collections .abc import Generator , Sequence
4+ from collections .abc import Generator
55from dataclasses import dataclass
6- from typing import Any
76
87
98# https://en.wikipedia.org/wiki/Tree_traversal
@@ -94,96 +93,86 @@ def height(root: Node | None) -> int:
9493 return (max (height (root .left ), height (root .right )) + 1 ) if root else 0
9594
9695
97- def level_order (root : Node | None ) -> Sequence [ Node | None ]:
96+ def level_order (root : Node | None ) -> Generator [ int , None , None ]:
9897 """
9998 Returns a list of nodes value from a whole binary tree in Level Order Traverse.
10099 Level Order traverse: Visit nodes of the tree level-by-level.
101100 """
102- output : list [Any ] = []
103101
104102 if root is None :
105- return output
103+ return
106104
107105 process_queue = deque ([root ])
108106
109107 while process_queue :
110108 node = process_queue .popleft ()
111- output . append ( node .data )
109+ yield node .data
112110
113111 if node .left :
114112 process_queue .append (node .left )
115113 if node .right :
116114 process_queue .append (node .right )
117- return output
118115
119116
120117def get_nodes_from_left_to_right (
121118 root : Node | None , level : int
122- ) -> Sequence [ Node | None ]:
119+ ) -> Generator [ int , None , None ]:
123120 """
124121 Returns a list of nodes value from a particular level:
125122 Left to right direction of the binary tree.
126123 """
127- output : list [Any ] = []
128124
129- def populate_output (root : Node | None , level : int ) -> None :
125+ def populate_output (root : Node | None , level : int ) -> Generator [ int , None , None ] :
130126 if not root :
131127 return
132128 if level == 1 :
133- output . append ( root .data )
129+ yield root .data
134130 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 )
137133
138- populate_output (root , level )
139- return output
134+ yield from populate_output (root , level )
140135
141136
142137def get_nodes_from_right_to_left (
143138 root : Node | None , level : int
144- ) -> Sequence [ Node | None ]:
139+ ) -> Generator [ int , None , None ]:
145140 """
146141 Returns a list of nodes value from a particular level:
147142 Right to left direction of the binary tree.
148143 """
149- output : list [Any ] = []
150144
151- def populate_output (root : Node | None , level : int ) -> None :
145+ def populate_output (root : Node | None , level : int ) -> Generator [ int , None , None ] :
152146 if root is None :
153147 return
154148 if level == 1 :
155- output . append ( root .data )
149+ yield root .data
156150 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 )
159153
160- populate_output (root , level )
161- return output
154+ yield from populate_output (root , level )
162155
163156
164- def zigzag (root : Node | None ) -> Sequence [ Node | None ] | list [ Any ]:
157+ def zigzag (root : Node | None ) -> Generator [ int , None , None ]:
165158 """
166159 ZigZag traverse:
167160 Returns a list of nodes value from left to right and right to left, alternatively.
168161 """
169162 if root is None :
170- return []
171-
172- output : list [Sequence [Node | None ]] = []
163+ return
173164
174165 flag = 0
175166 height_tree = height (root )
176167
177168 for h in range (1 , height_tree + 1 ):
178169 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 )
180171 flag = 1
181172 else :
182- output . append ( get_nodes_from_right_to_left (root , h ) )
173+ yield from get_nodes_from_right_to_left (root , h )
183174 flag = 0
184175
185- return output
186-
187176
188177def main () -> None : # Main function for testing.
189178 # Create binary tree.
@@ -198,15 +187,15 @@ def main() -> None: # Main function for testing.
198187 print (f"Height of Tree: { height (root )} " , "\n " )
199188
200189 print ("Complete Level Order Traversal: " )
201- print (level_order (root ), " \n " )
190+ print (f" { list ( level_order (root )) } \n " )
202191
203192 print ("Level-wise order Traversal: " )
204193
205194 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 ) ))
207196
208197 print ("\n ZigZag order Traversal: " )
209- print (zigzag (root ))
198+ print (f" { list ( zigzag (root )) } " )
210199
211200
212201if __name__ == "__main__" :
0 commit comments