diff --git a/data_structures/binary_tree/binary_tree_traversals.py b/data_structures/binary_tree/binary_tree_traversals.py index 5dbbbe623..2b33cdca4 100644 --- a/data_structures/binary_tree/binary_tree_traversals.py +++ b/data_structures/binary_tree/binary_tree_traversals.py @@ -1,9 +1,8 @@ from __future__ import annotations from collections import deque -from collections.abc import Generator, Sequence +from collections.abc import Generator from dataclasses import dataclass -from typing import Any # https://en.wikipedia.org/wiki/Tree_traversal @@ -94,96 +93,86 @@ def height(root: Node | None) -> int: return (max(height(root.left), height(root.right)) + 1) if root else 0 -def level_order(root: Node | None) -> Sequence[Node | None]: +def level_order(root: Node | None) -> Generator[int, None, None]: """ Returns a list of nodes value from a whole binary tree in Level Order Traverse. Level Order traverse: Visit nodes of the tree level-by-level. """ - output: list[Any] = [] if root is None: - return output + return process_queue = deque([root]) while process_queue: node = process_queue.popleft() - output.append(node.data) + yield node.data if node.left: process_queue.append(node.left) if node.right: process_queue.append(node.right) - return output def get_nodes_from_left_to_right( root: Node | None, level: int -) -> Sequence[Node | None]: +) -> Generator[int, None, None]: """ Returns a list of nodes value from a particular level: Left to right direction of the binary tree. """ - output: list[Any] = [] - def populate_output(root: Node | None, level: int) -> None: + def populate_output(root: Node | None, level: int) -> Generator[int, None, None]: if not root: return if level == 1: - output.append(root.data) + yield root.data elif level > 1: - populate_output(root.left, level - 1) - populate_output(root.right, level - 1) + yield from populate_output(root.left, level - 1) + yield from populate_output(root.right, level - 1) - populate_output(root, level) - return output + yield from populate_output(root, level) def get_nodes_from_right_to_left( root: Node | None, level: int -) -> Sequence[Node | None]: +) -> Generator[int, None, None]: """ Returns a list of nodes value from a particular level: Right to left direction of the binary tree. """ - output: list[Any] = [] - def populate_output(root: Node | None, level: int) -> None: + def populate_output(root: Node | None, level: int) -> Generator[int, None, None]: if root is None: return if level == 1: - output.append(root.data) + yield root.data elif level > 1: - populate_output(root.right, level - 1) - populate_output(root.left, level - 1) + yield from populate_output(root.right, level - 1) + yield from populate_output(root.left, level - 1) - populate_output(root, level) - return output + yield from populate_output(root, level) -def zigzag(root: Node | None) -> Sequence[Node | None] | list[Any]: +def zigzag(root: Node | None) -> Generator[int, None, None]: """ ZigZag traverse: Returns a list of nodes value from left to right and right to left, alternatively. """ if root is None: - return [] - - output: list[Sequence[Node | None]] = [] + return flag = 0 height_tree = height(root) for h in range(1, height_tree + 1): if not flag: - output.append(get_nodes_from_left_to_right(root, h)) + yield from get_nodes_from_left_to_right(root, h) flag = 1 else: - output.append(get_nodes_from_right_to_left(root, h)) + yield from get_nodes_from_right_to_left(root, h) flag = 0 - return output - def main() -> None: # Main function for testing. # Create binary tree. @@ -198,15 +187,15 @@ def main() -> None: # Main function for testing. print(f"Height of Tree: {height(root)}", "\n") print("Complete Level Order Traversal: ") - print(level_order(root), "\n") + print(f"{list(level_order(root))} \n") print("Level-wise order Traversal: ") for level in range(1, height(root) + 1): - print(f"Level {level}:", get_nodes_from_left_to_right(root, level=level)) + print(f"Level {level}:", list(get_nodes_from_left_to_right(root, level=level))) print("\nZigZag order Traversal: ") - print(zigzag(root)) + print(f"{list(zigzag(root))}") if __name__ == "__main__":