Invert binary tree in Python
Binary tree is a data structure and one of the simplest form of trees. You might have heard about people complaining  that during interviews they are asked to invert a binary tree. It may sound like something difficult, but in this article I’ll show you really simple solution using recursion (see this article for more recursion in Python). Inverting a tree basically means to switch places for right and left children of each node. Resulting tree will look like vertical mirroring of the input. Therefore if you know how to represent a tree within a code, you won’t stuck adding just a couple of extra lines invoking recursive function.
3.10is used throughout the article.
First of all we need a data structure which represents a tree and two helper functions:
generate_tree to create a target tree we plan to work with and
print_tree to visualize the result and verify our solution works as intended.
To represent a tree we need to define only one class that corresponds to a node. Each node stores some value/identificator as well as pointers to its children or
None in case of leaf nodes. An arbitrary variable assigned to the root node will state as a tree within our code.
Having this class defined we can go ahead and compose some simple tree by creating couple of linked nodes.
left_leaf = Node(23)
Storing one variable (root node) is enough to represent a whole tree as the rest of nodes are linked together using pointers.
Creating a tree manually is a hassle, so we need a function that generates arbitrary tree for us. We will provide number of levels as an argument and it will return a root node for the tree requested. Each node stores sequentually incremented value for better visual grasp, but you can also fill the tree with some random numbers.
from typing import Optional
We leverage recursion to generate left and right subtrees until we reach leaf nodes therefore returning
None for their children. To check whether generation was any good we need to add
print_tree function which is going to output its target to the console.
def print_tree(node: Node, level: int = 0):
Again, idea is to use recursion for traversing left and right subtrees and putting each node’s value to the terminal in between.
level parameter allows to add an extra indentation, so it’s visually clear on which level the node resides. Finally, we have a tree representation which looks like tree laying on its side (rotated counter-clockwise)
tree = generate_tree(3)
Use lines above if you want to produce same output as on the screenshot.
It’s no surprise that for inverting our tree we are going to use recursion again. This simple and straightforward solution requires even less code than generation itself.
def invert_tree(node: Node) -> Node:
The algorithm is the following: strarting from the root we invoke this function for the left and right subtrees and then swap them with each other. Therefore we end up having symmetrical tree from the same root node.
tree = generate_tree(3)
As you can see the resulting tree is symmetrical along the horizontal axis, so when folded on the dashed line corresponding nodes will match.
That’s basically it for the inversion itself. Clearly, there are more of extra code to help us represent and visualize the solution than within the solution itself. It gets tricky though when we want to accomplish the same without any recursion. Let’s move on to see how the same can be done using queue.
NOTE: There is also a slightly simpler solution using stack data structure. We are not going to implement it within the article as internally recursive solution works by storing all the function invocations on the stack. Basically that solution is equivalent of maintaining own call stack and essentially follows the exact same principle. Anyway, you can find source code for the stack-based solution within resources section in the end of the article.
The algorithm consists of two main steps: on the first stage we use breadth first search to traverse the tree and on the way we add leaf nodes to the intermediate list; on the second stage we restore tree-like structure from list elements rearraged in the desired order. Let’s look at each stage in more detail.
This is a simple implementation of a BFS that converts our input tree to the linear array of nodes.
def flatten_tree(node: Node) -> List[Node]:
The output of the function is a list containing all the elements in a specific order. This allows us to rebuild a tree attaching leaf nodes differently thus achieving requested order.
To understand better what this step does consider we are working with
[1, 5, 2, 7, 6, 4, 3] list as an input. You can obtain this exact order by going over our example tree from left to right column by column. Then based on the current position (
counter) we look ahead and mount leaf nodes back to the current node in the queue.
expand stages complement each other, so we should re-mount in the order opposite to the previous step. As a result left and right leaves for the each node got swapped.
def expand_into_tree(elems: List[Node]) -> Node:
There is also a helper
_get function simply to make getitem operation safe for our code. It makes sure no
IndexError occurs when we reach the end of our list and there is no more nodes to attach.
from typing import List, Optional, TypeVar
Code is a simple wrapper with
catch block returning
None when the index is out of bounds.
That’s everything we need to invert our tree non-recursively:
def invert_tree_queue(root: Node) -> Node:
Finally, confirm the solution works properly and matches previous results:
tree = generate_tree(3)
NOTE: Solution above works only with balanced trees where each non-leaf node has two children. This restriction comes up from the relying on the explicit order or the elements in the flattened tree. As an excersise you can modify the code to make it more general.
I don’t necessarily think that knowing all the algorithms and ability to write them on the whiteboard is a required thing for the programmer to have. As long as your code does something useful or cover some business need your are valuable for the field. Nowadays, especially with tools like GitHub Copilot and DeepMind AlphaCode you can write sophisticated code only by providing description of the feature required. Nevertheless, understanding main ideas and underlying concepts would hugely help you during development/debugging process. You don’t have to implement everything yourself or follow instructions step-by-step, but at least try to keep these articles as a side-reading. Stay curious, see ya.