Trie Data Structure: Autocomplete and Word Search in Python

Trie Data Structure

A Trie (prefix tree) stores strings character by character. It enables prefix searches in O(m) time (m = word length), making it perfect for autocomplete, spell checkers, and IP routing tables.

class TrieNode:
    def __init__(self):
        self.children = {}  # char -> TrieNode
        self.is_end   = False
        self.count    = 0   # how many words end here

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

    def insert(self, word: str) -> None:
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        node.is_end = True
        node.count += 1

    def search(self, word: str) -> bool:
        node = self._traverse(word)
        return node is not None and node.is_end

    def starts_with(self, prefix: str) -> bool:
        return self._traverse(prefix) is not None

    def _traverse(self, prefix: str):
        node = self.root
        for char in prefix:
            if char not in node.children:
                return None
            node = node.children[char]
        return node

trie = Trie()
for word in ['apple','app','application','apply','apt','banana']:
    trie.insert(word)

trie.search('app')         # True
trie.search('ap')          # False (not a complete word)
trie.starts_with('app')    # True  (prefix exists)

Autocomplete Feature

def autocomplete(trie, prefix: str, max_results: int = 10) -> list:
    node = trie._traverse(prefix)
    if not node: return []

    results = []

    def dfs(current_node, current_word):
        if len(results) >= max_results: return
        if current_node.is_end:
            results.append(current_word)
        for char, child in sorted(current_node.children.items()):
            dfs(child, current_word + char)

    dfs(node, prefix)
    return results

words = ['python', 'pytorch', 'pandas', 'php', 'perl', 'java', 'javascript']
trie2 = Trie()
for w in words: trie2.insert(w)

print(autocomplete(trie2, 'py'))  # ['python', 'pytorch']
print(autocomplete(trie2, 'p'))   # ['pandas', 'perl', 'php', 'python', 'pytorch']