Prefix Search Using Trie in DSA Python

Mental Model

A trie stores words by sharing common beginnings, so searching for words starting with a prefix is fast and easy.

Analogy: Imagine a tree where each branch is a letter, and all words that start the same way share the same path from the root.

root
 ↓
 a -> p -> p -> l -> e -> [end]
        ↓
        e -> [end]
 b -> a -> t -> [end]

Dry Run Walkthrough

Input: Insert words: apple, ape, bat; Search prefix: ap

Goal: Find all words starting with 'ap' quickly using the trie

Step 1: Insert 'apple' letter by letter into trie

root -> a -> p -> p -> l -> e -> [end]

Why: Build path for 'apple' so it can be found later

Step 2: Insert 'ape' sharing prefix 'ap' with 'apple'

root -> a -> p -> p -> l -> e -> [end]
           ↓
           e -> [end]

Why: Reuse existing 'ap' nodes to save space and time

Step 3: Insert 'bat' as a new branch

root -> a -> p -> p -> l -> e -> [end]
           ↓
           e -> [end]
 b -> a -> t -> [end]

Why: Add new word with different prefix

Step 4: Search prefix 'ap' by following nodes a -> p

At node 'p' after root -> a -> p

Why: Find the node where prefix ends to collect all words below

Step 5: Collect all words under 'ap' node: 'apple' and 'ape'

Words found: 'apple', 'ape'

Why: All words starting with 'ap' are under this node

Result:

Words starting with 'ap': apple, ape

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 search_prefix(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 below
        return results

# Driver code
trie = Trie()
words = ['apple', 'ape', 'bat']
for w in words:
    trie.insert(w)

prefix = 'ap'
found_words = trie.search_prefix(prefix)
print(f"Words starting with '{prefix}': {', '.join(found_words)}")

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

create new node for character if it does not exist

node = node.children[char] # move to child node

advance to next node along the word path

node.is_end = True # mark end of word

mark current node as end of a valid word

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

recursively collect all words under current node

for char in prefix: if char not in node.children: return [] # prefix not found node = node.children[char] # move to prefix node

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

Output

Words starting with 'ap': apple, ape

Complexity Analysis

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

Space: O(n) where n is total characters stored in trie nodes because each unique character creates a node

vs Alternative: Compared to scanning all words for prefix (O(n * m)), trie is faster for many searches as it shares prefix paths and avoids repeated checks

Edge Cases

Empty prefix search

Returns all words stored in trie

DSA Python

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

Prefix not present in trie

Returns empty list immediately

DSA Python

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

Insert duplicate words

Duplicates do not create extra nodes; word end marked again

DSA Python

node.is_end = True  # mark end of word

When to Use This Pattern

When you need to find all words starting with a prefix efficiently, use a trie because it stores shared prefixes once and allows fast prefix traversal.

Common Mistakes

Mistake: Not marking the end of a word, causing search to return incomplete results
Fix: Set node.is_end = True after inserting all characters of a word

Mistake: Not checking if prefix exists before DFS, leading to errors
Fix: Return empty list immediately if prefix path is missing in trie

Summary

Stores words by shared prefixes to enable fast prefix searches.

Use when you want to quickly find all words starting with a given prefix.

The key insight is that common prefixes share the same path, saving time and space.

Practice

1. What is the main purpose of using a Trie data structure in prefix search?

easy

A. To perform arithmetic operations quickly

B. To sort words alphabetically

C. To store numbers in sorted order

D. To efficiently find all words starting with a given prefix

Solution

1. Step 1: Understand Trie usage

   A Trie stores characters of words in a tree structure, allowing quick prefix lookups.

2. Step 2: Identify prefix search goal

   Prefix search means finding all words that start with a given prefix efficiently.

3. Final Answer:

   To efficiently find all words starting with a given prefix -> Option D

4. Quick Check:

   Trie = prefix search efficiency [OK]

Hint: Trie nodes represent prefixes for fast search [OK]

Common Mistakes:

• Confusing Trie with sorting data structure
• Thinking Trie is for arithmetic
• Assuming Trie stores whole words only

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

easy

A. node.children.add(char)

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

C. node.children.append(char)

D. node.children[char] = char

Solution

1. Step 1: Understand Trie children structure

   Trie children are usually stored in a dictionary mapping characters to TrieNode objects.

2. Step 2: Correct insertion syntax

   We check if character exists; if not, create new TrieNode and assign it to children dictionary.

3. Final Answer:

   node.children[char] = TrieNode() if char not in node.children else node.children[char] -> Option B

4. Quick Check:

   Insert word by mapping char to TrieNode [OK]

Hint: Use dictionary keys for children, not list methods [OK]

Common Mistakes:

• Using list methods like append or add on dictionary
• Assigning character instead of TrieNode
• Not checking if child node exists before insertion

3. Given the following code snippet, what will be the output of trie.starts_with('app')?

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 starts_with(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('apple')
trie.insert('app')
print(trie.starts_with('app'))

medium

A. Error

B. False

C. True

D. None

Solution

1. Step 1: Insert words 'apple' and 'app'

   Both words share prefix 'app', so nodes for 'a', 'p', 'p' exist in Trie.

2. Step 2: Check starts_with('app')

   Function checks each character in 'app' exists in children; all found, so returns True.

3. Final Answer:

   True -> Option C

4. Quick Check:

   Prefix 'app' exists in Trie = True [OK]

Hint: If prefix nodes exist, starts_with returns True [OK]

Common Mistakes:

• Confusing starts_with with full word search
• Expecting False because 'app' is prefix, not full word
• Ignoring that 'app' was inserted

4. Identify the bug in the following Trie starts_with method and select the correct fix:

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

medium

A. Return False inside else, return True after loop

B. Return True inside else, return True after loop

C. Return False inside else, return False after loop

D. Return True inside else, return False after loop

Solution

1. Step 1: Analyze current else return

   Currently, if char not found, method returns True incorrectly; should return False.

2. Step 2: Correct return values

   Return False inside else (prefix not found), and True after loop (prefix fully matched).

3. Final Answer:

   Return False inside else, return True after loop -> Option A

4. Quick Check:

   Prefix missing = False, prefix found = True [OK]

Hint: Return False if char missing, True if all found [OK]

Common Mistakes:

• Returning True when prefix char missing
• Returning False after loop instead of True
• Confusing else and after loop returns

5. You want to find all words in a Trie that start with the prefix "car". Which approach correctly extends the prefix search to list all such words?

hard

A. Traverse to node of 'car' prefix, then recursively collect all words from that node

B. Search each word in Trie and check if it starts with 'car' using string slicing

C. Insert 'car' as a word and then search for it

D. Use a hash map to store all words and filter by prefix

Solution

1. Step 1: Locate prefix node in Trie

   Traverse Trie nodes for 'c', 'a', 'r' to reach prefix node.

2. Step 2: Recursively collect words from prefix node

   From prefix node, explore all child nodes to build and collect complete words starting with 'car'.

3. Final Answer:

   Traverse to node of 'car' prefix, then recursively collect all words from that node -> Option A

4. Quick Check:

   Prefix node + recursion = all words with prefix [OK]

Hint: Find prefix node, then DFS to collect words [OK]

Common Mistakes:

• Checking all words instead of using Trie structure
• Inserting prefix as word incorrectly
• Using unrelated data structures like hash map