diff --git a/data_structures/binary_tree/binary_tree_mirror.py b/data_structures/binary_tree/binary_tree_mirror.py index b8548f4ec515..fa6628b6f728 100644 --- a/data_structures/binary_tree/binary_tree_mirror.py +++ b/data_structures/binary_tree/binary_tree_mirror.py @@ -1,45 +1,48 @@ -""" -Problem Description: -Given a binary tree, return its mirror. -""" - - -def binary_tree_mirror_dict(binary_tree_mirror_dictionary: dict, root: int): - if not root or root not in binary_tree_mirror_dictionary: +def mirror_subtree(tree: dict, root: int): + if root not in tree: return - left_child, right_child = binary_tree_mirror_dictionary[root][:2] - binary_tree_mirror_dictionary[root] = [right_child, left_child] - binary_tree_mirror_dict(binary_tree_mirror_dictionary, left_child) - binary_tree_mirror_dict(binary_tree_mirror_dictionary, right_child) + left, right = tree[root] + tree[root] = [right, left] + if left: + mirror_subtree(tree, left) + if right: + mirror_subtree(tree, right) -def binary_tree_mirror(binary_tree: dict, root: int = 1) -> dict: +def mirror_binary_tree(tree: dict, root: int = 1) -> dict: """ - >>> binary_tree_mirror({ 1: [2,3], 2: [4,5], 3: [6,7], 7: [8,9]}, 1) + Returns the mirror of the given binary tree starting from the root. + + >>> mirror_binary_tree({1: [2, 3], 2: [4, 5], 3: [6, 7], 7: [8, 9]}, 1) {1: [3, 2], 2: [5, 4], 3: [7, 6], 7: [9, 8]} - >>> binary_tree_mirror({ 1: [2,3], 2: [4,5], 3: [6,7], 4: [10,11]}, 1) + >>> mirror_binary_tree({1: [2, 3], 2: [4, 5], 3: [6, 7], 4: [10, 11]}, 1) {1: [3, 2], 2: [5, 4], 3: [7, 6], 4: [11, 10]} - >>> binary_tree_mirror({ 1: [2,3], 2: [4,5], 3: [6,7], 4: [10,11]}, 5) + >>> mirror_binary_tree({1: [2, 3], 2: [4, 5], 3: [6, 7], 4: [10, 11]}, 5) Traceback (most recent call last): ... - ValueError: root 5 is not present in the binary_tree - >>> binary_tree_mirror({}, 5) + ValueError: Root 5 is not present in the binary tree + >>> mirror_binary_tree({}, 5) Traceback (most recent call last): ... - ValueError: binary tree cannot be empty + ValueError: Binary tree cannot be empty """ - if not binary_tree: - raise ValueError("binary tree cannot be empty") - if root not in binary_tree: - msg = f"root {root} is not present in the binary_tree" - raise ValueError(msg) - binary_tree_mirror_dictionary = dict(binary_tree) - binary_tree_mirror_dict(binary_tree_mirror_dictionary, root) - return binary_tree_mirror_dictionary + if not tree: + raise ValueError("Binary tree cannot be empty") + if root not in tree: + error_message = f"Root {root} is not present in the binary tree" + raise ValueError(error_message) + + mirrored_tree = dict(tree) + mirror_subtree(mirrored_tree, root) + return mirrored_tree if __name__ == "__main__": + import doctest + + doctest.testmod() + binary_tree = {1: [2, 3], 2: [4, 5], 3: [6, 7], 7: [8, 9]} - print(f"Binary tree: {binary_tree}") - binary_tree_mirror_dictionary = binary_tree_mirror(binary_tree, 5) - print(f"Binary tree mirror: {binary_tree_mirror_dictionary}") + print(f"Original binary tree: {binary_tree}") + mirrored_tree = mirror_binary_tree(binary_tree, 1) + print(f"Mirrored binary tree: {mirrored_tree}")