A tree is a data structure that is normally used to represent hierarchical data. Python does not provide a  built-in implementation of the tree data structure.

Python user can choose to implement trees from scratch, or  use a third-party library. In this article we will explore the treelib library which  provides a simple and efficient implementation of the tree data structure.

treelib does not exist in the standard library, we therefore need to install it before we can start using it. You can run the following pip command to install the library:

pip install treelib

The Tree class

treelib defines the Tree class for representing trees. The following example shows a basic usage of the Tree class:

ExampleEdit & Run

create a tree 

#represents Abraham's Family tree

from treelib import Tree

#create a tree object
T = Tree()

T.create_node('Abraham', 'abraham')# The root node

#abraham's children
T.create_node('Ishmael', 'ishmael', parent = 'abraham')
T.create_node('Isaac', 'isaac', parent = 'abraham')
T.create_node('Zimran', 'zimran', parent = 'abraham')
T.create_node('Jockshan', 'jockshan', parent = 'abraham')
T.create_node('Midian', 'midian', parent = 'abraham')

#Ishmael's children
T.create_node('Nebaioth', 'nebaioth', parent = 'ishmael')
T.create_node('Kedar', 'kedar', parent = 'ishmael')
T.create_node('Tema', 'tema', parent = 'ishmael')

#Isaac's children
T.create_node('Esau', 'esau', parent = 'isaac')
T.create_node('Jacob', 'jacob', parent = 'isaac')

#Esau's children
T.create_node('Eiphaz', 'eiphaz', parent = 'esau')
T.create_node('Reuel', 'reuel', parent = 'esau')
T.create_node('Korah', 'korah', parent = 'esau')

#jacob's children
T.create_node('Reuben', 'reuben', parent = 'jacob')
T.create_node('Joseph', 'joseph', parent = 'jacob')
T.create_node('Levi', 'levi', parent = 'jacob')
T.create_node('Simeon', 'simeon', parent = 'jacob')
T.create_node('Benjamin', 'benjamin', parent = 'jacob')

#display the tree
print(T)
Output:
Traceback (most recent call last):  File "<string>", line 3, in <module>ModuleNotFoundError: No module named 'treelib'[Finished in 0.011252525029703975s]

The above example shows an incomplete representations of Abraham's family tree.

As you can see from the above example, the Tree class provides an intuitive API for creating tree nodes and establishing their relationships.

At the heart of the Tree class is the create_node() method which creates a node and establishes its position in the hierarchy. The method has the following syntax:

create_node(tag = None, identifier = None, parent = None, data = None)
tag A human readable string for the object stored in the node. It defaults to None.
identifier A unique identifier to identify this node in the tree. If not given, a random unique value will be auto-generated.
parent The identifier of this node's parent. It defaults to None.
data The object stored in the node. It defaults to None.

Tree operations

In this section we will look at the various operations that can be performed on Tree objects.

Get the height of the tree

To get the height of the tree, use the depth() method. If called without any argument, this method returns the height of the entire tree, whereas with argument, it will return the depth of the node whose identifier is given as argument.

ExampleEdit & Run
#The tree is omitted refer to the previous sections

print(T.depth())
print(T.depth('jacob')) #get the depth of node 'jacob'
Output:

3
2

Get a node using its identifier

Use the get_node() method to get a reference to the node object whose identifier is given as argument.

Syntax:
get_node(nid)

Where nid is the identifier of the node.

ExampleEdit & Run
#The tree is omitted refer to the previous sections

n = T.get_node('joseph')
print(n)
print(n.tag)
Output:

Node(tag=Joseph, identifier=joseph, data=None)
Joseph

Get the subtree rooted at a given node

The subtree() method returns the subtree rooted at the node whose identifier is given as argument.

Syntax:
subtree(nid)
ExampleEdit & Run
#The tree is omitted refer to the previous sections

subtree = T.subtree('isaac')
print(subtree)
Output:

Isaac
├── Esau
│   ├── Eiphaz
│   ├── Korah
│   └── Reuel
└── Jacob
    ├── Benjamin
    ├── Joseph
    ├── Levi
    ├── Reuben
    └── Simeon

Remove the subtree rooted at a given node

use remove_subtree() method to remove the subtree that is rooted at the node whose identifier is given as argument.

Syntax:
remove_subtree(nid)

The method removes and returns the tree rooted at the node whose identifier is nid. If nid is not given, the method returns an empty tree. 

ExampleEdit & Run
#The tree is omitted refer to the previous sections

T.remove_subtree('jacob')
print(T)
Output:

Abraham
├── Isaac
│   └── Esau
│       ├── Eiphaz
│       ├── Korah
│       └── Reuel
├── Ishmael
│   ├── Kedar
│   ├── Nebaioth
│   └── Tema
├── Jockshan
├── Midian
└── Zimran

Attach another tree to a given node

The paste() method attaches a new tree as a subtree of the node whose identifier is given in arguments.

Syntax:
paste(nid, new_tree)
ExampleEdit & Run
#The tree is omitted refer to the previous sections

new_tree = Tree()
new_tree.create_node("n1", "n1")
new_tree.create_node('n2', 'n2', parent = 'n1')
new_tree.create_node('n3', 'n3', parent = 'n2')

#attach the tree to the original tree
T.paste('joseph', new_tree)
Output:

 Abraham
├── Isaac
│   ├── Esau
│   │   ├── Eiphaz
│   │   ├── Korah
│   │   └── Reuel
│   └── Jacob
│       ├── Benjamin
│       ├── Joseph
│       │   └── n1
│       │       └── n2
│       │           └── n3
│       ├── Levi
│       ├── Reuben
│       └── Simeon
├── Ishmael
│   ├── Kedar
│   ├── Nebaioth
│   └── Tema
├── Jockshan
├── Midian
└── Zimran

View All Tree Methods

Use the dir() function to view all the methods that are associated with Tree objects.  You can use the builtin filter() function to exclude dunder methods from the returned list.

ExampleEdit & Run
from treelib import Tree

print(*filter(lambda x: '__' not in x and not x.isupper(), dir(Tree)), sep = "\n")
Output:
Traceback (most recent call last):  File "<string>", line 1, in <module>ModuleNotFoundError: No module named 'treelib'[Finished in 0.011203200090676546s]

Use the builtin help() function with a method to view its usage and other helpful information.

ExampleEdit & Run
from treelib import Tree

help(Tree.all_nodes)
Output:
Traceback (most recent call last):  File "<string>", line 1, in <module>ModuleNotFoundError: No module named 'treelib'[Finished in 0.010708692949265242s]