Autocomplete System with Trie in DSA Python

Mental Model

A trie stores words by sharing common prefixes, so we can quickly find all words starting with a given prefix.

Analogy: Imagine a tree of roads where each road sign is a letter. To find all shops starting with 'ca', you follow the roads 'c' then 'a' and see all shops branching from there.

root
 ├─ c
 │   ├─ a
 │   │   ├─ t (end)
 │   │   └─ r (end)
 │   └─ o
 │       └─ w (end)
 └─ d
     └─ o
         └─ g (end)

Dry Run Walkthrough

Input: words: ['cat', 'car', 'cow', 'dog'], prefix: 'ca'

Goal: Find all words starting with 'ca' using the trie

Step 1: Insert 'cat' into the trie letter by letter

root -> c -> a -> t[end]

Why: We add each letter as a node, marking 't' as end of word

Step 2: Insert 'car' sharing prefix 'ca', add 'r' node

root -> c -> a -> t[end], r[end]

Why: Reuse 'c' and 'a' nodes, add 'r' as new branch

Step 3: Insert 'cow', add 'o' and 'w' nodes after 'c'

root -> c -> a -> t[end], r[end]; c -> o -> w[end]

Why: Add new branch for 'cow' starting with 'co'

Step 4: Insert 'dog', add 'd', 'o', 'g' nodes from root

root -> c -> a -> t[end], r[end]; c -> o -> w[end]; d -> o -> g[end]

Why: Add new branch for 'dog' starting with 'd'

Step 5: Search prefix 'ca' by following 'c' then 'a'

At node 'a' with children 't[end]', 'r[end]'

Why: We reached prefix node to find all words starting with 'ca'

Step 6: Collect all words under 'a' node: 'cat' and 'car'

Found words: ['cat', 'car']

Why: These are all words starting with prefix 'ca'

Result:

root -> c -> a -> t[end], r[end]; c -> o -> w[end]; d -> o -> g[end]
Words with prefix 'ca': ['cat', 'car']

Annotated Code

DSA Python

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

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()  # create node if missing
            node = node.children[char]  # move to child node
        node.is_end = True  # mark end of word

    def _dfs(self, node: TrieNode, prefix: str, results: list) -> None:
        if node.is_end:
            results.append(prefix)  # found a word
        for char, child in node.children.items():
            self._dfs(child, prefix + char, results)  # explore deeper

    def autocomplete(self, prefix: str) -> list:
        node = self.root
        for char in prefix:
            if char not in node.children:
                return []  # prefix not found
            node = node.children[char]  # move to prefix node
        results = []
        self._dfs(node, prefix, results)  # collect all words
        return results

# Driver code
words = ['cat', 'car', 'cow', 'dog']
trie = Trie()
for w in words:
    trie.insert(w)

prefix = 'ca'
results = trie.autocomplete(prefix)
print(f"Words with prefix '{prefix}': {results}")

if char not in node.children: node.children[char] = TrieNode()

create new node if letter path missing

node = node.children[char]

advance to child node to continue insertion or search

node.is_end = True

mark current node as end of a valid word

for char, child in node.children.items(): self._dfs(child, prefix + char, results)

recursively collect all words starting from prefix node

for char in prefix: if char not in node.children: return [] node = node.children[char]

traverse trie nodes to reach prefix end or return empty if missing

Output

Words with prefix 'ca': ['cat', 'car']

Complexity Analysis

Time: O(m + k) where m is length of prefix and k is total characters in matched words because we traverse prefix then collect all matches

Space: O(n * l) where n is number of words and l is average word length for storing trie nodes

vs Alternative: Compared to scanning all words for prefix match O(n*l), trie reduces prefix search to O(m) plus output size, making autocomplete efficient

Edge Cases

Empty prefix

Returns all words in the trie

DSA Python

for char in prefix:
    if char not in node.children:
        return []

Prefix not in trie

Returns empty list immediately

DSA Python

if char not in node.children:
    return []

Words with shared prefixes

Trie shares nodes, no duplication, all correct words returned

DSA Python

node.children[char] = TrieNode()

When to Use This Pattern

When you need to find all words starting with a prefix efficiently, use a trie to quickly navigate shared prefixes and collect matches.

Common Mistakes

Mistake: Not marking the end of a word, causing autocomplete to miss valid words
Fix: Set node.is_end = True after inserting all letters of a word

Mistake: Returning partial matches without checking is_end, including prefixes that are not words
Fix: Only add words to results when node.is_end is True during DFS

Summary

Stores words in a tree structure to quickly find all words starting with a given prefix.

Use when you want fast autocomplete or prefix search among many words.

The key is sharing common prefix nodes so you don't repeat work for each word.

Practice

1. What is the main purpose of using a Trie in an autocomplete system?

easy

A. To sort words alphabetically

B. To efficiently store and search prefixes of words

C. To count the total number of words

D. To encrypt the stored words

Solution

1. Step 1: Understand the role of Trie in prefix search

   A Trie is a tree-like data structure designed to store strings by their prefixes, allowing quick prefix-based searches.

2. Step 2: Connect Trie usage to autocomplete

   Autocomplete requires finding all words starting with a given prefix efficiently, which Trie supports by traversing nodes matching the prefix.

3. Final Answer:

   To efficiently store and search prefixes of words -> Option B

4. Quick Check:

   Trie = prefix search [OK]

Hint: Trie stores prefixes for fast autocomplete search [OK]

Common Mistakes:

• Thinking Trie sorts words alphabetically
• Confusing Trie with hash tables
• Assuming Trie encrypts data

2. Which of the following is the correct way to insert a word "cat" into a Trie node in Python?

easy

A. node['c'] = TrieNode()

B. node.insert('c')

C. node.add('c')

D. node.children['c'] = TrieNode()

Solution

1. Step 1: Understand Trie node structure

   Trie nodes typically have a dictionary called children to map characters to child nodes.

2. Step 2: Insert character as a key in children dictionary

   To insert 'c', assign a new TrieNode to node.children['c'].

3. Final Answer:

   node.children['c'] = TrieNode() -> Option D

4. Quick Check:

   Insert char in children dict = node.children['c'] = TrieNode() [OK]

Hint: Use children dictionary to add new character nodes [OK]

Common Mistakes:

• Trying to assign directly to node['c'] without children
• Using non-existent methods like add or insert
• Confusing list with dictionary for children

3. Given the following code snippet, what will be the output?

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

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 = True

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

trie = Trie()
trie.insert('dog')
trie.insert('deer')
print(trie.search_prefix('de'))
print(trie.search_prefix('do'))
print(trie.search_prefix('da'))

medium

A. True\nFalse\nFalse

B. False\nTrue\nFalse

C. True\nTrue\nFalse

D. False\nFalse\nFalse

Solution

1. Step 1: Insert words 'dog' and 'deer' into Trie

   Both words add nodes for their characters. 'dog' adds d->o->g, 'deer' adds d->e->e->r.

2. Step 2: Check prefixes 'de', 'do', and 'da'

   'de' exists as prefix in 'deer', so returns True. 'do' exists as prefix in 'dog', so True. 'da' does not exist, so False.

3. Final Answer:

   True True False -> Option C

4. Quick Check:

   Prefix 'de' and 'do' found, 'da' not found = True, True, False [OK]

Hint: Check prefix path existence in Trie children [OK]

Common Mistakes:

• Confusing search_prefix with full word search
• Assuming 'da' prefix exists
• Ignoring is_end flag for prefix search

4. The following code is intended to insert words into a Trie, but it raises an error. What is the error?

class TrieNode:
    def __init__(self):
        self.children = []
        self.is_end = False

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]

medium

A. TypeError: list indices must be integers or slices, not str

B. AttributeError: 'TrieNode' object has no attribute 'children'

C. SyntaxError: invalid syntax

D. No error, code runs fine

Solution

1. Step 1: Identify data structure used for children

   The code uses a list for children, but tries to access it with string keys.

2. Step 2: Understand list indexing rules

   Lists require integer indices, but code uses node.children[char] where char is a string, causing TypeError.

3. Final Answer:

   TypeError: list indices must be integers or slices, not str -> Option A

4. Quick Check:

   Using list with string keys causes TypeError [OK]

Hint: Use dictionary, not list, for children to index by characters [OK]

Common Mistakes:

• Using list instead of dict for children
• Ignoring error message details
• Assuming code runs without error

5. You want to implement an autocomplete system that returns all words starting with a given prefix using a Trie. Which approach correctly collects all matching words after reaching the prefix node?

hard

A. Perform a depth-first search (DFS) from the prefix node, collecting words where is_end is True

B. Return all children keys of the prefix node without further traversal

C. Only return the prefix itself without searching deeper

D. Use breadth-first search (BFS) but ignore is_end flags

Solution

1. Step 1: Reach the node matching the prefix

   Traverse the Trie to the node representing the last character of the prefix.

2. Step 2: Collect all words starting from prefix node using DFS

   Use DFS to explore all child nodes recursively, collecting words where is_end is True to get complete words.

3. Final Answer:

   Perform a depth-first search (DFS) from the prefix node, collecting words where is_end is True -> Option A

4. Quick Check:

   DFS from prefix node with is_end check = correct autocomplete [OK]

Hint: Use DFS from prefix node, collect words at is_end True [OK]

Common Mistakes:

• Returning only immediate children keys
• Ignoring is_end flags and returning incomplete words
• Using BFS without checking word ends