Implementing a Trie Data Structure in Python

Introduction

A Trie, pronounced as “try,” is a tree-like data structure that is used to store an associative array where the keys are usually strings.

Unlike binary trees, which have a maximum of two children (left and right), tries can have multiple children, one for each letter in the alphabet.

Tries are particularly useful for searches in dictionaries, providing auto-suggestions during text searches, and even IP routing.

Structure of a Trie

A trie typically starts with an empty root node, and each path down the tree represents a single word. The nodes themselves are not inherently meaningful; it’s the path to the node from the root that defines the key or word.

Below is a representation of a trie storing “bat,” “batman,” “batmobile,” “batwoman,” and “batgirl”:

      root
        |
        b
        |
        a
        |
        t--m--a--n
        |   |
        |   w--o--m--a--n
        |   |
        |   g--i--r--l
        |
        m--o--b--i--l--e

Python Implementation

Let’s implement a basic trie to store words. We’ll also add methods to insert words, search for words, and even to suggest completions based on partial words.

Trie Node Class

Firstly, let’s define a node in the trie.

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end_of_word = False

Trie Class

Now we can implement the Trie class, which will use TrieNode objects.

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

    def insert(self, word):
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        node.is_end_of_word = True

    def search(self, word):
        node = self.root
        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]
        return node.is_end_of_word

    def start_with_prefix(self, prefix):
        node = self.root
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        return True

    def suggestions(self, node, word, prefix):
        if node.is_end_of_word:
            print(prefix + word)

        for key, child_node in node.children.items():
            self.suggestions(child_node, word + key, prefix)

    def print_suggestions(self, prefix):
        node = self.root
        for char in prefix:
            if char not in node.children:
                return
            node = node.children[char]
        self.suggestions(node, "", prefix)

Method Details:

• insert: Adds a word to the trie.
• search: Checks if a whole word exists in the trie.
• start_with_prefix: Checks if there are any words in the trie that start with the given prefix.
• suggestions: Internal function to print all words in the trie that starts with a given prefix.
• print_suggestions: Initiates the process of printing all words starting with a given prefix.

Using the Trie Class

Here’s a quick example of how to use the Trie class.

# Create a new Trie
trie = Trie()

# Insert words
trie.insert("bat")
trie.insert("batman")
trie.insert("batmobile")
trie.insert("batwoman")
trie.insert("batgirl")

# Search for a word
print(trie.search("bat"))  # Output: True
print(trie.search("batwing"))  # Output: False

# Check prefix
print(trie.start_with_prefix("batw"))  # Output: True

# Print suggestions
trie.print_suggestions("bat")

The suggestions will output:

bat
batman
batmobile
batwoman
batgirl

If you find this guide helpful, please share with others, and spread the knowledge!