treelib in Python

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:

#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)

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)

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.

#The tree is omitted refer to the previous sections

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

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.

get_node(nid)

Where nid is the identifier of the node.

#The tree is omitted refer to the previous sections

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

Get the subtree rooted at a given node

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

subtree(nid)

#The tree is omitted refer to the previous sections

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

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.

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. 

#The tree is omitted refer to the previous sections

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

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.

paste(nid, new_tree)

#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)

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.

from treelib import Tree

print(*filter(lambda x: '__' not in x and not x.isupper(), dir(Tree)), sep = "\n")

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

from treelib import Tree

help(Tree.all_nodes)