Binary Search Trees, Traversals and Balancing in Python

Data Structures and Algorithms in Python is beginner-friendly introduction to common data structures (linked lists, stacks, queues, graphs) and algorithms (search, sorting, recursion, dynamic programming) in Python, designed to help you prepare for coding interviews and assessments.

This tutorial is an executable Jupyter notebook. Click the Open in Colab button at the top of this page to execute the code.

Jupyter Notebooks: This notebook is made of cells. Each cell can contain code written in Python or explanations in plain English. You can execute code cells and view the results instantly within the notebook. Jupyter is a powerful platform for experimentation and analysis. Don't be afraid to mess around with the code & break things - you'll learn a lot by encountering and fixing errors. You can use the "Kernel > Restart & Clear Output" menu option to clear all outputs and start again from the top.

Problem

In this notebook, we'll focus on solving the following problem:

QUESTION 1: As a senior backend engineer at Jovian, you are tasked with developing a fast in-memory data structure to manage profile information (username, name and email) for 100 million users. It should allow the following operations to be performed efficiently:

1. Insert the profile information for a new user.
2. Find the profile information of a user, given their username
3. Update the profile information of a user, given their usrname
4. List all the users of the platform, sorted by username

You can assume that usernames are unique.

Along the way, we will also solve several other questions related to binary trees and binary search trees that are often asked in coding interviews and assessments.

The Method

Here's a systematic strategy we'll apply for solving problems:

1. State the problem clearly. Identify the input & output formats.
2. Come up with some example inputs & outputs. Try to cover all edge cases.
3. Come up with a correct solution for the problem. State it in plain English.
4. Implement the solution and test it using example inputs. Fix bugs, if any.
5. Analyze the algorithm's complexity and identify inefficiencies, if any.
6. Apply the right technique to overcome the inefficiency. Repeat steps 3 to 6.

1. State the problem clearly. Identify the input & output formats.

Problem

We need to create a data structure which can store 100 million records and perform insertion, search, update and list operations efficiently.

Input

The key inputs to our data structure are user profiles, which contain the username, name and email of a user.

A Python class would be a great way to represent the information for a user. A class is a blueprint for creating objects. Everything in Python is an object belonging to some class. Here's the simples possible class in Python, with nothing in it:

We can create or instantiate an object of the class by calling it like a function.

We can verify that the object is of the class User.

The object user1 does not contain any useful information. Let's add a constructor method to the class to store some attributes or properties.

We can now create an object with some properties.

Here's what's happening above (conceptually):

• Python creates an empty object of the type user and stores in the variable user2
• Python then invokes the function User.___init__ with the arguments user2, "john", "John Doe" and "john@doe.com"
• As the __init__ function is executed, the properties username, name and email are set on the object user2

We can access the properties of the object using the . notation.

You can also define custom methods inside a class.

When we try to invoke the method user3.introduce_yourself, the object user3 is automatically passed as the first argument self. Indeed, the following statement is equivalent to the above statement.

Finally, we'll define a couple of helper methods to display user objects nicely within Jupyter.

Exercise: What is the purpose of defining the functions __str__ and __repr__ within a class? How are the two functions different? Illustrate with some examples using the empty cells below.

Output

We can also express our desired data structure as a Python class UserDatabase with four methods: insert, find, update and list_all.

It's good programming practice to list out the signatures of different class functions before we actually implement the class.

2. Come up with some example inputs & outputs.

Let's create some sample user profiles that we can use to test our functions once we implement them.

We can access different fields within a user profile object using the . (dot) notation.

We can also view a string representation of the object, since defined the __repr__ and __str__ methods

Since we haven't implemented our data structure yet, it's not possible to list sample outputs. However you can try to come up with different scenarios to test future implementations

Exercise: List some scenarios for testing the class methods insert, find, update and list_all.

1. Insert:

   1. Inserting into an empty database of users
   2. Trying to insert a user with a username that already exists
   3. Inserting a user with a username that does not exist
   4. ???

2. Find:

   1. ???
   2. ???
   3. ???

3. Update:

   1. ???
   2. ???
   3. ???

4. List:

   1. ???
   2. ???
   3. ???

3. Come up with a correct solution. State it in plain English.

Here's a simple and easy solution to the problem: we store the User objects in a list sorted by usernames.

The various functions can be implemented as follows:

1. Insert: Loop through the list and add the new user at a position that keeps the list sorted.
2. Find: Loop through the list and find the user object with the username matching the query.
3. Update: Loop through the list, find the user object matching the query and update the details
4. List: Return the list of user objects.

We can use the fact usernames, which are are strings can be compared using the <, > and == operators in Python.

4. Implement the solution and test it using example inputs.

The code for implementing the above solution is also fairly straightfoward.

We can create a new database of users by instantiating and object of the UserDatabase class.

Let's insert some entires into the object.

We can now retrieve the data for a user, given their username.

Let's try changing the information for a user

Finally, we can retrieve a list of users in alphabetical order.

Let's verify that a new user is inserted into the correct position.

The user biraj was inserted just before hemanth, as expected.

Exercise: Use the empty cells below to test the various scenarios you listed in step 2 above.

5. Analyze the algorithm's complexity and identify inefficiencies

The operations insert, find, update involves iterating over a list of users, in the worst case, they may take up to N iterations to return a result, where N is the total number of users. list_all however, simply returns the existing internal list of users.

Thus, the time complexities of the various operations are:

1. Insert: O(N)
2. Find: O(N)
3. Update: O(N)
4. List: O(1)

Exercise: Verify that the space complexity of each operation is O(1).

Is this good enough? To get a sense how long each function might take if there are 100 million users on the platform, we can simply run an for or while loop on 10 million numbers.

It takes almost 10 seconds to execute all the iterations in the above cell.

• A 10-second delay for fetching user profiles will lead to a suboptimal users experience and may cause many users to stop using the platform altogether.
• The 10-second processing time for each profile request will also significantly limit the number of users that can access the platform at a time or increase the cloud infrastructure costs for the company by millions of dollars.

As a senior backend engineer, you must come up with a more efficient data structure! Choosing the right data structure for the requirements at hand is an important skill. It's apparent that a sorted list of users might not be the best data structure to organize profile information for millions of users.

6. Apply the right technique to overcome the inefficiency

We can limit the number of iterations required for common operations like find, insert and update by organizing our data in the following structure, called a binary tree:

It's called a tree because it vaguely like an inverted tree trunk with branches.

• The word "binary" indicates that each "node" in the tree can have at most 2 children (left or right).
• Nodes can have 0, 1 or 2 children. Nodes that do not have any children are sometimes also called "leaves".
• The single node at the top is called the "root" node, and it typically where operations like search, insertion etc. begin.

Balanced Binary Search Trees

For our use case, we require the binary tree to have some additional properties:

1. Keys and Values: Each node of the tree stores a key (a username) and a value (a User object). Only keys are shown in the picture above for brevity. A binary tree where nodes have both a key and a value is often referred to as a map or treemap (because it maps keys to values).
2. Binary Search Tree: The left subtree of any node only contains nodes with keys that are lexicographically smaller than the node's key, and the right subtree of any node only contains nodes with keys that lexicographically larger than the node's key. A tree that satisfies this property is called a binary search trees, and it's easy to locate a specific key by traversing a single path down from the root note.
3. Balanced Tree: The tree is balanced i.e. it does not skew too heavily to one side or the other. The left and right subtrees of any node shouldn't differ in height/depth by more than 1 level.

Height of a Binary Tree

The number of levels in a tree is called its height. As you can tell from the picture above, each level of a tree contains twice as many nodes as the previous level.

For a tree of height k, here's a list of the number of nodes at each level:

Level 0: 1

Level 1: 2

Level 2: 4 i.e. 2^2

Level 3: 8 i.e. 2^3

...

Level k-1: 2^(k-1)

If the total number of nodes in the tree is N, then it follows that

N = 1 + 2^1 + 2^2 + 2^3 + ... + 2^(k-1)

We can simplify this equation by adding 1 on each side:

N + 1 = 1 + 1 + 2^1 + 2^2 + 2^3 + ... + 2^(k-1) 

N + 1 = 2^1 + 2^1 + 2^2+ 2^3 + ... + 2^(k-1) 

N + 1 = = 2^2 + 2^2 + 2^3 + ... + 2^(k-1)

N + 1 = = 2^3 + 2^3 + ... + 2^(k-1)

...

N + 1 = 2^(k-1) + 2^(k-1)

N + 1 = 2^k

k = log(N + 1) <= log(N) + 1 

Thus, to store N records we require a balanced binary search tree (BST) of height no larger than log(N) + 1. This is a very useful property, in combination with the fact that nodes are arranged in a way that makes it easy to find a specific key by following a single path down from the root.

As we'll see soon, the insert, find and update operations in a balanced BST have time complexity O(log N) since they all involve traversing a single path down from the root of the tree.

Binary Tree

QUESTION 2: Implement a binary tree using Python, and show its usage with some examples.

To begin, we'll create simple binary tree (without any of the additional properties) containing numbers as keys within nodes. Here's an example:

Here's a simple class representing a node within a binary tree.

Let's create objects representing each node of the above tree

Let's verify that node0 is an object of the type TreeNode and has the property key set to 3.

We can connect the nodes by setting the .left and .right properties of the root node.

And we're done! We can create a new variable tree which simply points to the root node, and use it to access all the nodes within the tree.

Going forward, we'll use the term "tree" to refer to the root node. The term "node" can refer to any node in a tree, not necessarily the root.

Exercise: Create the following binary tree using the TreeNode class defined above.

It's a bit inconvenient to create a tree by manually connecting all the nodes. Let's write a helper function which can convert a tuple with the structure ( left_subtree, key, right_subtree) (where left_subtree and right_subtree are themselves tuples) into binary tree.

Here's an tuple representing the tree shown above:

The parse_tuple creates a new root node when a tuple of size 3 as an the input. Interestingly, to create the left and right subtrees for the node, the parse_tuple function invokes itself. This technique is called recursion. The chain of recursive calls ends when parse_tuple encounters a number or None as input. We'll use recursion extensively throughout this tutorial.

Exercise: Add print statements inside parse_tuple to display the arguments for each call of the function. Does the sequence of recursive calls make sense to you?

Let's try out parse_tuple with the tuple define earlier.

We can now examine the tree to verify that it was constructed as expected.

Exercise: Define a function tree_to_tuple that converts a binary tree into a tuple representing the same tree. E.g. tree_to_tuple converts the tree created above to the tuple ((1, 3, None), 2, ((None, 3, 4), 5, (6, 7, 8))). Hint: Use recursion.

Let's create another helper function to display all the keys in a tree-like structure for easier visualization.

Once again, the display_keys function users recursion to print all the keys of the left and right subtree with proper indentation.

Exercise: Add print statements inside display_keys to display the arguments for each call of the function. Does the sequence of recursive calls make sense to you?

Let's try using the function.

We can now visualize the tree that was just created (albeit rotated by 90 degrees). It's easy to see that it matches the expected structure.

Exercise: Create some more trees and visualize them using display_keys. You can use excalidraw.com as a digital whiteboard to create trees.

Traversing a Binary Tree

The following questions are frequently asked in coding interviews and assessments:

QUESTION 3: Write a function to perform the inorder traversal of a binary tree.

QUESTION 4: Write a function to perform the preorder traversal of a binary tree.

QUESTION 5: Write a function to perform the postorder traversal of a binary tree.

A traversal refers to the process of visiting each node of a tree exactly once. Visiting a node generally refers to adding the node's key to a list. There are three ways to traverse a binary tree and return the list of visited keys:

Inorder traversal

1. Traverse the left subtree recursively inorder.
2. Traverse the current node.
3. Traverse the right subtree recursively inorder.

Preorder traversal

1. Traverse the current node.
2. Traverse the left subtree recursively preorder.
3. Traverse the right subtree recursively preorder.

Can you guess how postorder traversal works??

Here's an implementation of inorder traversal of a binary tree.

Let's try it out with this tree:

Exercise: Implement functions for preorder and postorder traversal of a binary tree.

Test your implementations by making submissions to the following problems:

• https://leetcode.com/problems/binary-tree-inorder-traversal/
• https://leetcode.com/problems/binary-tree-preorder-traversal/
• https://leetcode.com/problems/binary-tree-postorder-traversal/

Height and Size of a Binary Tree

QUESTION 6: Write a function to calculate the height/depth of a binary tree

QUESTION 7: Write a function to count the number of nodes in a binary tree

The height/depth of a binary tree is defined as the length of the longest path from its root node to a leaf. It can be computed recursively, as follows:

Let's compute the height of this tree:

Here's a function to count the number of nodes in a binary tree.

Exercise: Try solving these problems relating to path lengths in a binary tree:

• https://leetcode.com/problems/maximum-depth-of-binary-tree/
• https://leetcode.com/problems/minimum-depth-of-binary-tree/
• https://leetcode.com/problems/diameter-of-binary-tree/

As a final step, let's compile all the functions we've written so far as methods withing the TreeNode class itself. Encapsulation of data and functionality within the same class is a good programming practice.

The class method invocations TreeNode.height(node) and node.height() are equivalent. Can you guess why we're using the former in the function definitions above? Let's try out the various methods defined above for this tree:

Exercise: Create some more trees and try out the operations defined above. Add more operations to the TreeNode class.

Binary Search Tree (BST)

A binary search tree or BST is a binary tree that satisfies the following conditions:

1. The left subtree of any node only contains nodes with keys less than the node's key
2. The right subtree of any node only contains nodes with keys greater than the node's key

It follows from the above conditions that every subtree of a binary search tree must also be a binary search tree.

QUESTION 8: Write a function to check if a binary tree is a binary search tree (BST).

QUESTION 9: Write a function to find the maximum key in a binary tree.

QUESTION 10: Write a function to find the minimum key in a binary tree.

Here's a function that covers all of the above:

The following tree is not a BST (because a node with the key 3 appears in the left subtree of a node with the key 2):

Let's verify this using is_bst.

On the other hand, the following tree is a BST:

Let's create this tree and verify that it is a BST. Note that the TreeNode class also supports using strings as keys, as strings support the comparison operators < and > too.

Exercise: Test the is_bst function with some more examples using the empty cells below.

Storing Key-Value Pairs using BSTs

Recall that we need to store user objects with each key in our BST. Let's define new class BSTNode to represent the nodes of of our tree. Apart from having properties key, left and right, we'll also store a value and pointer to the parent node (for easier upward traversal).

Let's try to recreate this BST with usernames as keys and user objects as values:

Exercise: Add the next layer of nodes to the tree and verify that they were added properly.

We can use the same display_keys function we defined earlier to visualize our tree.

Insertion into BST

QUESTION 11: Write a function to insert a new node into a BST.

We use the BST-property to perform insertion efficiently:

1. Starting from the root node, we compare the key to be inserted with the current node's key
2. If the key is smaller, we recursively insert it in the left subtree (if it exists) or attach it as as the left child if no left subtree exists.
3. If the key is larger, we recursively insert it in the right subtree (if it exists) or attach it as as the right child if no right subtree exists.

Here's a recursive implementation of insert.

Let's use this to recreate our tree.

To create the first node, we can use the insert function with None as the target tree.

The remaining nodes can now be inserted into tree.

Perfect! The tree was created as expected.

Note, however, that the order of insertion of nodes change the structure of the resulting tree.

Can you see why the tree created above is skewed/unbalanced?

Skewed/unbalanced BSTs are problematic because the height of such trees often ceases to logarithmic compared to the number of nodes in the tree. For instance the above tree has 7 nodes and height 7.

The length of the path traversed by insert is equal to the height of the tree (in the worst case). It follows that if the tree is balanced, the time complexity of insertion is O(log N) otherwise it is O(N).

Exercise: Create some more balanced and unbalanced BSTs using the insert function defined above.

Finding a Node in BST

QUESTION 11: Find the value associated with a given key in a BST.

We can follow a recursive strategy similar to insertion to find the node with a given key within a BST.

The the length of the path followed by find is equal to the height of the tree (in the worst case). Thus it has a similar time complexity as insert.

Example: Try finding some more nodes from the BST created above (or create new trees).

Updating a value in a BST

QUESTION 12: Write a function to update the value associated with a given key within a BST

We can use find to locate the node to be updated, and simply update it's value.

The value of the node was successfully updated. The time complexity of update is the same as that of find.

Exercise: Try some more update operations using the BST created earlier.

List the nodes

QUESTION 13: Write a function to retrieve all the key-values pairs stored in a BST in the sorted order of keys.

The nodes can be listed in sorted order by performing an inorder traversal of the BST.

Exercise: Determine the time complexity and space complexity of list_all.

Balanced Binary Trees

QUESTION 14: Write a function to determine if a binary tree is balanced.

Here's a recursive strategy:

1. Ensure that the left subtree is balanced.
2. Ensure that the right subtree is balanced.
3. Ensure that the difference between heights of left subtree and right subtree is not more than 1.

The following tree is balanced:

The following tree is not balanced:

Exercise: Is the tree shown below balanced? Why or why not? Create this tree and check if it's balanced using the is_balanced function.

Also try this related problem on complete binary trees: https://leetcode.com/problems/check-completeness-of-a-binary-tree/

Balanced Binary Search Trees

QUESTION 15: Write a function to create a balanced BST from a sorted list/array of key-value pairs.

We can use a recursive strategy here, turning the middle element of the list into the root, and recursively creating left and right subtrees.

Recall that the same list of users, when inserted one-by-one resulted in a skewed tree.

Balancing an Unbalanced BST

QUESTION 16: Write a function to balance an unbalanced binary search tree.

We first perform an inorder traversal, then create a balanced BST using the function defined earlier.

After every insertion, we can balance the tree. This way the tree will remain balanced.

Complexity of the various operations in a balanced BST:

• Insert - O(log N) + O(N) = O(N)
• Find - O(log N)
• Update - O(log N)
• List all - O(N)

What's the real improvement between O(N) and O(log N)?

The logarithm (base 2) of 100 million is around 26. Thus, it takes only 26 operations to find or update a node within a BST (as opposed to 100 million).

Compared to linear time:

Thus, find and update from a balanced binary search tree is 300,000 times faster than our original solution. To speed up insertions, we may choose to perform the balancing periodically (e.g. once every 1000 insertions). This way, most insertions will be O (log N), but every 1000th insertion will take a few seconds. Another options is to rebalance the tree periodically at the end of every hour.

A Python-Friendly Treemap

We are now ready to return to our original problem statement.

QUESTION 1: As a senior backend engineer at Jovian, you are tasked with developing a fast in-memory data structure to manage profile information (username, name and email) for 100 million users. It should allow the following operations to be performed efficiently:

1. Insert the profile information for a new user.
2. Find the profile information of a user, given their username
3. Update the profile information of a user, given their usrname
4. List all the users of the platform, sorted by username

You can assume that usernames are unique.

We can create a generic class TreeMap which supports all the operations specified in the original problem statement in a python-friendly manner.

Exercise: What is the time complexity of __len__? Can you reduce it to O(1). Hint: Modify the BSTNode class.

Let's try using the TreeMap class below.

Exercise: Try out some more examples below. Can our treemap actually handle millions of users profiles?

Self-Balancing Binary Trees and AVL Trees

A self-balancing binary tree remains balanced after every insertion or deletion. Several decades of research has gone into creating self-balancing binary trees, and many approaches have been devised e.g. B-trees, Red Black Trees and AVL (Adelson-Velsky Landis) trees.

We'll take a brief look at AVL trees. Self-balancing in AVL trees is achieved by tracking the balance factor (difference between the height of the left subtree and the right subtree) for each node and rotating unbalanced subtrees along the path of insertion/deletion to balance them.

In a balanced BST, the balance factor of each node is either 0, -1, or 1. When we perform an insertion, then the balance factor of certain nodes along the path of insertion may change to 2 or -2. Those nodes can be "rotated" one-by-one to bring the balance factor back to 1, 0 or -1.

There are 4 different scenarios for balancing, two of which require a single rotation, while the others require 2 rotations:

Source: HackerRank

Since each rotation takes constant time, and at most log N rotations may be required, this operation is far more efficient than creating a balanced binary tree from scratch, allowing insertion and deletion to be performed in O (log N) time. Here are some references for AVL Trees:

• Explanation of the various cases: https://youtu.be/jDM6_TnYIqE?t=482
• Implementation: https://www.geeksforgeeks.org/avl-tree-set-1-insertion/

Summary and Exercises

Binary trees form the basis of many modern programming language features (e.g. maps in C++ and Java) and data storage systems (filesystem indexes, relational databases like MySQL). You might wonder if dictionaries in Python are also binary search trees. They're not. They're hash tables, which is a different but equally interesting and important data structure. We'll explore hash tables in a future tutorial.

We've covered a lot of ground this in this tutorial, including several common interview questions. Here are a few more problems you can try out:

1. Implement rotations and self-balancing insertion
2. Implement deletion of a node from a binary search tree
3. Implement deletion of a node from a BST (with balancing)
4. Find the lowest common ancestor of two nodes in a tree (Hint: Use the parent property)
5. Find the next node in lexicographic order for a given node
6. Given a number k, find the k-th node in a BST.

Try more questions here:

• https://medium.com/techie-delight/binary-tree-interview-questions-and-practice-problems-439df7e5ea1f
• https://leetcode.com/tag/tree/