Fixed Errors. Code was not working

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}")