Trees in Python

Now that we knows the basics of a tree data structure, we can start to think about how we would actually use one in Python.

Just as we saw with linked lists, there is also no built-in tree data structure in Python. There are tree libraries that you can import, but we won't do so in this class. Instead, it's instructive for us to build a tree class from first principles.

The LinkedTree class

We will create a LinkedTree class to represent a binary tree, where each node in the tree has a value and two possible child pointers: a left and a right. To do so, we will need our tree to be composed of tree nodes, which we will create using a nested, inner class:

class LinkedTree:
    class TreeNode:
        def __init__(self, val, left=None, right=None):
            self.val = val
            self.left = left
            self.right = right

    def __init__(self, root):
        self.root = root

Notice that each TreeNode will have a value, a (possibly None) left pointer, and a (possibly None) right pointer.

Building a LinkedTree

Just like with linked lists, we could create some TreeNode objects, link them together manually, and use one of them as the root of a new LinkedTree. For example:

n1 = TreeNode(5)
n2 = TreeNode(10)
n3 = TreeNode(7)
n1.left = n2
n1.right = n3
tree = LinkedTree(n1)

After constructing this tree, this is what the picture would look like in memory:

Note that we represent each TreeNode as a box with three parts: a left child pointer, a value, and a right child pointer.

Building a tree with random values

We could also write a method to construct a random tree for us. Watch the video below to see how to build of height h filled with random values:

Counting the number of nodes in a tree

With our LinkedTree class, we can now write algorithms that operate on the tree. See the below implementation of a num_nodes() method, which finds the number of nodes in a tree:

def __num_nodes(self, node):
    if node is None:
        return 0

    return 1 + self.__num_nodes(node.left) + self.__num_nodes(node.right)

def num_nodes(self):
    return self.__num_nodes(self.root)

The basic algorithm is that the count of nodes in the tree is equivalent to 1 (for the root node) + the number of nodes in the left subtree + the number of nodes in the subtree. When applied recursively, this will count the total number of nodes in the tree.

Note: It is much easier to write recursive algorithms on trees than iterative algorithms, since we need to be able to traverse up and down the tree. Making recursive calls and then returning from those recursive calls models this up and down traversal naturally. Therefore, we generally won't write iterative algorithms on trees, with some exceptions.

As always, we are interested in finding the time complexity of our algorithm. For num_nodes(), we make one recursive call for each node in the tree. As the number of nodes in the tree increases, the number of recursive calls increases linearly. Therefore, the running time of num_nodes() is O(n).

Finding the depth of a node

Watch the video below to see how to write a method that finds the depth of a node: