Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Mirror a Binary Tree solution #9534

153 changes: 153 additions & 0 deletions data_structures/binary_tree/mirror_binary_tree.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
"""
Given the root of a binary tree, mirror the tree, and return its root.

Leetcode problem reference: https://leetcode.com/problems/mirror-binary-tree/
"""
from __future__ import annotations

from collections.abc import Iterator
from dataclasses import dataclass


@dataclass
class Node:
"""
A Node has value variable and pointers to Nodes to its left and right.
"""

value: int
left: Node | None = None
right: Node | None = None

def __iter__(self) -> Iterator[int]:
if self.left:
yield from self.left
yield self.value
if self.right:
yield from self.right

def __len__(self) -> int:
return sum(1 for _ in self)

def mirror(self) -> Node:
"""
Mirror the binary tree rooted at this node by swapping left and right children.

>>> tree = Node(0)
>>> list(tree)
[0]
>>> list(tree.mirror())
[0]
>>> tree = Node(1, Node(0), Node(3, Node(2), Node(4, None, Node(5))))
>>> tuple(tree)
(0, 1, 2, 3, 4, 5)
>>> tuple(tree.mirror())
(5, 4, 3, 2, 1, 0)
"""
self.left, self.right = self.right, self.left
if self.left:
self.left.mirror()
if self.right:
self.right.mirror()
return self


def make_tree_seven() -> Node:
r"""
Return a binary tree with 7 nodes that looks like this:
1
/ \
2 3
/ \ / \
4 5 6 7

>>> tree_seven = make_tree_seven()
>>> len(tree_seven)
7
>>> list(tree_seven)
[4, 2, 5, 1, 6, 3, 7]
"""
tree = Node(1)
tree.left = Node(2)
tree.right = Node(3)
tree.left.left = Node(4)
tree.left.right = Node(5)
tree.right.left = Node(6)
tree.right.right = Node(7)
return tree


def make_tree_nine() -> Node:
r"""
Return a binary tree with 9 nodes that looks like this:
1
/ \
2 3
/ \ \
4 5 6
/ \ \
7 8 9

>>> tree_nine = make_tree_nine()
>>> len(tree_nine)
9
>>> list(tree_nine)
[7, 4, 8, 2, 5, 9, 1, 3, 6]
"""
tree = Node(1)
tree.left = Node(2)
tree.right = Node(3)
tree.left.left = Node(4)
tree.left.right = Node(5)
tree.right.right = Node(6)
tree.left.left.left = Node(7)
tree.left.left.right = Node(8)
tree.left.right.right = Node(9)
return tree


def main() -> None:
r"""
Mirror binary trees with the given root and returns the root

>>> tree = make_tree_nine()
>>> tuple(tree)
(7, 4, 8, 2, 5, 9, 1, 3, 6)
>>> tuple(tree.mirror())
(6, 3, 1, 9, 5, 2, 8, 4, 7)

nine_tree:
1
/ \
2 3
/ \ \
4 5 6
/ \ \
7 8 9

The mirrored tree looks like this:
1
/ \
3 2
/ / \
6 5 4
/ / \
9 8 7
"""
trees = {"zero": Node(0), "seven": make_tree_seven(), "nine": make_tree_nine()}
for name, tree in trees.items():
print(f" The {name} tree: {tuple(tree)}")
# (0,)
# (4, 2, 5, 1, 6, 3, 7)
# (7, 4, 8, 2, 5, 9, 1, 3, 6)
print(f"Mirror of {name} tree: {tuple(tree.mirror())}")
# (0,)
# (7, 3, 6, 1, 5, 2, 4)
# (6, 3, 1, 9, 5, 2, 8, 4, 7)


if __name__ == "__main__":
import doctest

doctest.testmod()
main()