How to Implement Trie Data Structure in Python

1. Trie Data Structure in Python
2. Insert a String in a Trie in Python
3. Search an Element in a Trie in Python
4. Trie Implementation in Python

Tries are the data structures that we use to store strings. They allow us to search text strings in the most efficient manner possible.

This article will discuss how we can implement a Trie in Python.

Trie Data Structure in Python

You can consider a Trie as a tree where each node consists of a character. Each node has one or more children depending on whether it is an internal character of a string or the last character.

The node representing the last character of a string has no child, and it marks the end of the string. We will include a flag variable in the class definition to mark the end of a string.

Each node in the Trie can have a maximum of 26 children if we store strings that consist of only small case English letters. If the strings have characters other than the alphabets, the maximum number of children of a particular node equals the total number of distinct characters.

You can define a Node class in Python to implement a node of a Trie, as shown in the following example.

class Node:
    def __init__(self):
        self.children = []
        for i in range(26):
            self.children.append(None)
        self.isLeafNode = False

Here, we created a list named children used to define whether or not a character is a child of the present node. As we considered 26 characters, we initialized the list with 26 None values.

If a character is not a child of the current node, its position will contain the value None. Otherwise, the position corresponding to that character stores the node for that character.

While inserting characters into the children list, we store the children nodes in the alphabetical order of the characters. In other words, the child node of the letter a will be stored at index 0, the child node of the letter b will be stored at index 1, etc.

After creating a node, we need to create a class to define a Trie. In the class, we will define an empty node with a list containing 26 None values to represent the 26 characters of the English alphabet.

We will call the empty Node the root node.

class Trie:
    def __init__(self):
        self.root = Node()

Whenever a string is inserted into the Trie, the node representing the string’s first character becomes the child of the root node. Note that we will store the nodes containing the next characters of the strings as list elements according to their position.

After creating the root node, we will implement the methods to insert a word in the Trie and search for a word in the Trie in the following sections.

Insert a String in a Trie in Python

To insert a character into the Trie, we will first find the length of the string that is to be inserted. After that, we will start crawling the Trie from the root node of the Trie.

The following is the algorithm to insert a string into the Trie:

1. Calculate the length of the string to be inserted into the Trie. Store it in a variable strLen.
2. Take a variable crawler and assign the root node of the Trie to the variable.
3. If you are at level n, check if the nth character of the string is present at that level in the Trie. If yes, store its position in the children list in a variable position; then, go to 5 - otherwise, go to 4.
4. Create a new node for the Trie and assign it to the index position of the crawler.
5. Move the crawler to the next level.
6. Check if we have reached the end of the string; if yes, go to 7 - otherwise, go to 3.
7. Mark the current node as the end of the string.

Having discussed the algorithm, let us now implement this algorithm to insert a string into a Trie in Python.

def insert(self, input_str):
    strLen = len(input_str)
    crawler = self.root
    for level in range(strLen):
        character = input_str[level]
        position = ord(character) - ord("a")
        if crawler.children[position] is None:
            crawler.children[position] = Node()
        crawler = crawler.children[position]
    crawler.isLeafNode = True

Search an Element in a Trie in Python

To search whether a string is present in a Trie or not, we will use the following algorithm.

1. Initialize a variable crawler and assign the root node of the Trie to the variable.
2. Calculate the length of the string to be searched in the Trie. Store it in a variable strLen.
3. At a level n, find if the nth character of the string is present in the children list. If yes, go to 4; otherwise, return False.
4. Check if the current node is a leaf node. If yes, return True; else, increment n and go to 3.

We have defined the algorithm for searching a string in a Trie. Let us implement it in Python.

def search(self, input_str):
    crawler = self.root
    strLen = len(input_str)
    for level in range(strLen):
        character = input_str[level]
        position = ord(character) - ord("a")
        if crawler.children[position] is None:
            return False
        crawler = crawler.children[position]
    return crawler.isLeafNode

Trie Implementation in Python

As we have implemented the methods for search and insert operations in a Trie in Python, let us execute the code using some example operations.

class Node:
    def __init__(self):
        self.children = []
        for i in range(26):
            self.children.append(None)
        self.isLeafNode = False


class Trie:
    def __init__(self):
        self.root = Node()

    def insert(self, input_str):
        strLen = len(input_str)
        crawler = self.roothave
        for level in range(strLen):
            character = input_str[level]
            position = ord(character) - ord("a")
            if crawler.children[position] is None:
                crawler.children[position] = Node()
            crawler = crawler.children[position]
        crawler.isLeafNode = True

    def search(self, input_str):
        crawler = self.root
        strLen = len(input_str)
        for level in range(strLen):
            character = input_str[level]
            position = ord(character) - ord("a")
            if crawler.children[position] is None:
                return False
            crawler = crawler.children[position]
        return crawler.isLeafNode


x = Trie()
myStr = "aditya"
print("Inserting the string:", myStr)
x.insert(myStr)
myStr = "delftstack"
print("Inserting the string:", myStr)
x.insert(myStr)
myStr = "aaditya"
print("Inserting the string:", myStr)
x.insert(myStr)
print("aditya is present in the trie:", x.search("aditya"))
print("delftstack is present in the trie:", x.search("delftstack"))
print("python is present in the trie:", x.search("python"))

Output:

Inserting the string: aditya
Inserting the string: delftstack
Inserting the string: aaditya
aditya is present in the trie: True
delftstack is present in the trie: True
python is present in the trie: False

We first implemented a Trie in Python using the algorithms discussed above in this example. After that, we inserted three strings, aditya, delftstack, and aaditya into the Trie.

Then, we performed search operations on the Trie to check if the strings aditya, delftstack, and python were present in the Trie or not. You can observe the output in the example.