Skip to content

Instantly share code, notes, and snippets.

@auvipy
Forked from tonybaloney/InvertBinaryTree.py
Last active August 27, 2020 07:11
Show Gist options
  • Save auvipy/fa1aee41c0c2d27e2cab0d86e23475e4 to your computer and use it in GitHub Desktop.
Save auvipy/fa1aee41c0c2d27e2cab0d86e23475e4 to your computer and use it in GitHub Desktop.
Invert binary tree in Python
from dataclasses import dataclass
from typing import Any
from xml.etree.ElementTree import Element
@dataclass
class Node():
val: Any
right: "Node"
left: "Node"
def __str__(self): # DFS recursive-print
return " ".join((str(self.left), str(self.val), str(self.right))).replace("None", "")
"""
1
/ \
2 3
/ \ / \
4 5 6 7
"""
# build tree
tree = Node(1, None, None)
tree.left = Node(2, None, None)
tree.right = Node(3, None, None)
tree.left.left = Node(4, None, None)
tree.left.right = Node(5, None, None)
tree.right.left = Node(6, None, None)
tree.right.right = Node(7, None, None)
print(tree)
# Option 1 (sensible): recursively swap left and right properties
def invert(root: Node):
if not root:
return
root.left, root.right = invert(root.right), invert(root.left)
return root
# Option 2 The Enterprise way:
def enterpriseInvert(root: Node):
import xml.etree.ElementTree as ET
def btreeToElementTree(root: Node, el: Element):
el.attrib['val'] = str(root.val)
if root.left:
left = ET.SubElement(el, 'left')
btreeToElementTree(root.left, left)
if root.right:
right = ET.SubElement(el, 'right')
btreeToElementTree(root.right, right)
return el
tree = btreeToElementTree(root, ET.Element('tree'))
# Swap left and right via XPath searches
lefts = tree.findall(".//left")
rights = tree.findall(".//right")
for left in lefts:
left.tag = "right"
for right in rights:
right.tag = "left"
return ET.tostring(tree)
print(invert(tree))
print(enterpriseInvert(tree))
"""
Prints
4 2 5 1 6 3 7
7 3 6 1 5 2 4
b'<tree val="1"><right val="3"><right val="7" /><left val="6" /></right><left val="2"><right val="5" /><left val="4" /></left></tree>'
"""
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment