Implement Trie (Prefix Tree) in Python

A Trie (prefix tree) is a tree data structure used to efficiently store and search strings. It's particularly useful for autocomplete features and prefix matching. Each node represents a character, and paths from root to nodes represent prefixes.

Trie Structure Overview

A Trie supports three main operations ?

• insert(word) − Adds a word to the trie
• search(word) − Returns True if the complete word exists
• startsWith(prefix) − Returns True if any word starts with the given prefix

Implementation

We'll use a dictionary-based approach where each node contains child dictionaries and an end marker ?

class Trie:
    def __init__(self):
        """Initialize the trie with an empty root node."""
        self.child = {}
    
    def insert(self, word):
        """Insert a word into the trie."""
        current = self.child
        for letter in word:
            if letter not in current:
                current[letter] = {}
            current = current[letter]
        current['#'] = True  # Mark end of word
    
    def search(self, word):
        """Search for a complete word in the trie."""
        current = self.child
        for letter in word:
            if letter not in current:
                return False
            current = current[letter]
        return '#' in current
    
    def startsWith(self, prefix):
        """Check if any word starts with the given prefix."""
        current = self.child
        for letter in prefix:
            if letter not in current:
                return False
            current = current[letter]
        return True

# Example usage
trie = Trie()
trie.insert("apple")
print(trie.search("apple"))     # Complete word exists
print(trie.search("app"))       # Incomplete word
print(trie.startsWith("app"))   # Prefix exists
trie.insert("app")
print(trie.search("app"))       # Now complete word exists

The output of the above code is ?

True
False
True
True

How It Works

The Trie uses nested dictionaries to represent the tree structure ?

• Root node − Empty dictionary representing the starting point
• Character nodes − Each key represents a character, value is another dictionary
• End marker − '#' key indicates a complete word ends at this node

Insertion Process

When inserting "apple" ?

1. Start at root dictionary
2. Create path: 'a' ? 'p' ? 'p' ? 'l' ? 'e'
3. Add '#' marker at the final 'e' node

Advanced Example

Let's build a more comprehensive example with multiple words ?

class Trie:
    def __init__(self):
        self.child = {}
    
    def insert(self, word):
        current = self.child
        for letter in word:
            if letter not in current:
                current[letter] = {}
            current = current[letter]
        current['#'] = True
    
    def search(self, word):
        current = self.child
        for letter in word:
            if letter not in current:
                return False
            current = current[letter]
        return '#' in current
    
    def startsWith(self, prefix):
        current = self.child
        for letter in prefix:
            if letter not in current:
                return False
            current = current[letter]
        return True
    
    def getAllWords(self, prefix=""):
        """Get all words with given prefix."""
        words = []
        current = self.child
        
        # Navigate to prefix node
        for letter in prefix:
            if letter not in current:
                return words
            current = current[letter]
        
        # Find all words from this node
        def dfs(node, word):
            if '#' in node:
                words.append(word)
            for char, child in node.items():
                if char != '#':
                    dfs(child, word + char)
        
        dfs(current, prefix)
        return words

# Building a word dictionary
trie = Trie()
words = ["cat", "car", "card", "care", "careful", "dog", "dodge"]

for word in words:
    trie.insert(word)

print("Words starting with 'car':", trie.getAllWords("car"))
print("Does 'care' exist?", trie.search("care"))
print("Any words start with 'do'?", trie.startsWith("do"))

The output shows the trie's prefix matching capabilities ?

Words starting with 'car': ['car', 'card', 'care', 'careful']
Does 'care' exist? True
Any words start with 'do'? True

Time Complexity

Where m is the length of the word or prefix being processed.

Conclusion

Tries provide efficient string storage and retrieval with O(m) time complexity for all operations. They're ideal for autocomplete systems, spell checkers, and IP routing tables where prefix matching is crucial.