Prefix Search Using Trie in DSA Go

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. Words that start the same way share the same path down the tree, like roads sharing the same route before splitting.

root
 ↓
 t -> r -> i -> e -> null
 ↑
 prefix search starts here

Dry Run Walkthrough

Input: Insert words: "tree", "trie", "algo", "algo"; Search prefix: "tr"

Goal: Find all words starting with prefix "tr"

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

root -> t -> r -> e -> e -> null

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

Step 2: Insert 'trie' sharing 't' and 'r', then add 'i' and 'e'

root -> t -> r -> e -> e -> null
           ↓
           i -> e -> null

Why: Reuse existing prefix 'tr' to save space and time

Step 3: Insert 'algo' as new branch from root

root -> a -> l -> g -> o -> null
  ↓
  t -> r -> e -> e -> null
           ↓
           i -> e -> null

Why: Add new word starting with different letter

Step 4: Search prefix 'tr' by following 't' then 'r'

At node 'r' with children 'e' and 'i'

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

Step 5: Collect all words below 'r': 'tree' and 'trie'

Words found: 'tree', 'trie'

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

Result:

Words with prefix 'tr': tree, trie

Annotated Code

DSA Go

package main

import (
	"fmt"
)

// TrieNode represents each node in the trie
// It holds children nodes and a flag if it's end of a word
type TrieNode struct {
	children map[rune]*TrieNode
	isEnd    bool
}

// Trie structure with root node
type Trie struct {
	root *TrieNode
}

// NewTrieNode creates a new TrieNode
func NewTrieNode() *TrieNode {
	return &TrieNode{children: make(map[rune]*TrieNode)}
}

// NewTrie creates a new Trie with root node
func NewTrie() *Trie {
	return &Trie{root: NewTrieNode()}
}

// Insert adds a word into the trie
func (t *Trie) Insert(word string) {
	curr := t.root
	for _, ch := range word {
		if curr.children[ch] == nil {
			curr.children[ch] = NewTrieNode() // create node if missing
		}
		curr = curr.children[ch] // move to child node
	}
	curr.isEnd = true // mark end of word
}

// searchNode returns the node where prefix ends or nil if not found
func (t *Trie) searchNode(prefix string) *TrieNode {
	curr := t.root
	for _, ch := range prefix {
		if curr.children[ch] == nil {
			return nil // prefix not found
		}
		curr = curr.children[ch] // move to next node
	}
	return curr
}

// collectWords recursively collects all words below given node
func (t *Trie) collectWords(node *TrieNode, prefix string, results *[]string) {
	if node.isEnd {
		*results = append(*results, prefix) // found a word
	}
	for ch, child := range node.children {
		t.collectWords(child, prefix+string(ch), results) // explore children
	}
}

// StartsWith returns all words starting with given prefix
func (t *Trie) StartsWith(prefix string) []string {
	results := []string{}
	node := t.searchNode(prefix) // find node for prefix
	if node == nil {
		return results // no words with prefix
	}
	t.collectWords(node, prefix, &results) // collect all words below
	return results
}

func main() {
	trie := NewTrie()
	words := []string{"tree", "trie", "algo", "algo"}
	for _, w := range words {
		trie.Insert(w)
	}
	prefix := "tr"
	results := trie.StartsWith(prefix)
	fmt.Printf("Words with prefix '%s': %s\n", prefix, results)
}

if curr.children[ch] == nil { curr.children[ch] = NewTrieNode() // create node if missing }

create new node if letter path does not exist

curr = curr.children[ch] // move to child node

advance current node to next letter node

curr.isEnd = true // mark end of word

mark node as end of a valid word

if curr.children[ch] == nil { return nil // prefix not found }

stop search if prefix letter missing

if node.isEnd { *results = append(*results, prefix) // found a word }

add word to results when end of word reached

t.collectWords(child, prefix+string(ch), results) // explore children

recursively collect all words below current node

Output

Words with prefix 'tr': [tree trie]

Complexity Analysis

Time: O(m + k) because inserting or searching a word takes time proportional to its length m, and collecting k words with the prefix takes additional time

Space: O(n * m) because in worst case all words are different and stored separately, each of length m

vs Alternative: Compared to scanning all words for prefix (O(n*m)), trie is faster for prefix search as it avoids checking every word

Edge Cases

Empty prefix search

Returns all words stored in trie

DSA Go

if node == nil {
	return results // no words with prefix
}

Prefix not present in trie

Returns empty list

DSA Go

if curr.children[ch] == nil {
	return nil // prefix not found
}

Duplicate words inserted

Duplicates do not affect trie structure or results

DSA Go

curr.isEnd = true // mark end of word

When to Use This Pattern

When you need to quickly find all words starting with a prefix, use a trie because it shares common prefixes and speeds up search.

Common Mistakes

Mistake: Not marking the end of word node, causing prefix search to miss valid words
Fix: Always set isEnd = true at the last node of inserted word

Summary

Stores words in a tree structure sharing prefixes for fast prefix search.

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

The key is that common prefixes share the same path, so searching is quick and space-saving.

Practice

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

easy

A. To sort words alphabetically

B. To efficiently find all words that start with a given prefix

C. To store numbers in sorted order

D. To perform arithmetic operations quickly

Solution

1. Step 1: Understand what prefix search means

   Prefix search means finding all words that begin with a certain starting sequence of letters.

2. Step 2: Understand Trie's role in prefix search

   A Trie stores words in a tree-like structure where each node represents a character, making it fast to find all words sharing a prefix.

3. Final Answer:

   To efficiently find all words that start with a given prefix -> Option B

4. Quick Check:

   Prefix search = efficient word lookup [OK]

Hint: Trie nodes represent characters for fast prefix lookup [OK]

Common Mistakes:

• Confusing Trie with sorting data structures
• Thinking Trie is for arithmetic operations
• Assuming Trie stores numbers only

2. Which of the following is the correct way to declare a Trie node struct in Go for prefix search?

easy

A. type TrieNode struct { value int; left *TrieNode; right *TrieNode }

B. type TrieNode struct { children []int; isEnd int }

C. type TrieNode struct { data string; next *TrieNode }

D. type TrieNode struct { children map[rune]*TrieNode; isEnd bool }

Solution

1. Step 1: Identify the Trie node components

   A Trie node should have a map from characters (rune) to child nodes and a boolean to mark word end.

2. Step 2: Match the correct Go struct syntax

   type TrieNode struct { children map[rune]*TrieNode; isEnd bool } uses a map[rune]*TrieNode and a bool isEnd, which fits the Trie node definition.

3. Final Answer:

   type TrieNode struct { children map[rune]*TrieNode; isEnd bool } -> Option D

4. Quick Check:

   Trie node = map + bool [OK]

Hint: Trie nodes use map[rune]*TrieNode and bool isEnd [OK]

Common Mistakes:

• Using slice of int instead of map for children
• Confusing linked list node with Trie node
• Using int instead of bool for isEnd

3. Given the following Go code snippet for inserting words into a Trie and searching for prefix "app":

trie := NewTrie()
trie.Insert("apple")
trie.Insert("app")
trie.Insert("application")
result := trie.StartsWith("app")
fmt.Println(result)

What will be the output of fmt.Println(result)?

medium

A. [app]

B. [apple application]

C. [apple app application]

D. []

Solution

1. Step 1: Insert words into the Trie

   Words "apple", "app", and "application" are inserted, all starting with prefix "app".

2. Step 2: Search for prefix "app"

   The StartsWith method returns all words starting with "app", which includes all three inserted words.

3. Final Answer:

   [apple app application] -> Option C

4. Quick Check:

   All words with prefix "app" = [apple app application] [OK]

Hint: StartsWith returns all words sharing the prefix [OK]

Common Mistakes:

• Missing the exact prefix word "app" in results
• Returning only one matching word
• Returning empty list if prefix matches

4. Consider this buggy Go code snippet for searching a prefix in a Trie:

func (t *Trie) StartsWith(prefix string) []string {
    node := t.root
    for _, ch := range prefix {
        if node.children[ch] == nil {
            return nil
        }
        node = node.children[ch]
    }
    return node.words
}

What is the main bug causing incorrect prefix search results?

medium

A. The code does not collect all words under the prefix node

B. The code does not check if node.words is initialized

C. The code returns nil instead of empty slice when prefix not found

D. The code uses wrong variable name for children map

Solution

1. Step 1: Analyze the prefix traversal loop

   The loop correctly moves node down the Trie for each character in prefix.

2. Step 2: Check the return statement after traversal

   Returning node.words assumes all words under prefix are stored there, but usually words are stored in leaf nodes, so it misses collecting all words under subtree.

3. Final Answer:

   The code does not collect all words under the prefix node -> Option A

4. Quick Check:

   Prefix search must gather all subtree words [OK]

Hint: Collect all subtree words after prefix node for correct results [OK]

Common Mistakes:

• Returning nil instead of empty slice is minor, not main bug
• Assuming words stored only at prefix node
• Ignoring recursive collection of words

5. You want to implement a prefix search that returns only words starting with a prefix but excludes words that are exactly the prefix itself. Which modification to the Trie search logic achieves this?

hard

A. After finding prefix node, return all words except the one equal to prefix

B. Return only words that have prefix length or fewer characters

C. Skip marking isEnd for prefix nodes during insertion

D. Return all words including prefix itself

Solution

1. Step 1: Understand the requirement

   We want all words starting with prefix but exclude the word exactly equal to prefix.

2. Step 2: Modify search result filtering

   After collecting all words under prefix node, filter out the word that matches prefix exactly before returning.

3. Final Answer:

   After finding prefix node, return all words except the one equal to prefix -> Option A

4. Quick Check:

   Filter out exact prefix word from results [OK]

Hint: Filter out exact prefix word after collecting all matches [OK]

Common Mistakes:

• Removing isEnd flag breaks normal word insertion
• Returning only longer words by length may miss valid words
• Returning prefix itself defeats the requirement