Autocomplete System with Trie in DSA Go

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 places starting with 'ca', you follow the roads 'c' then 'a' and see all destinations branching from there.

root
 ↓
 c -> a -> t [end]
       ↓
       r -> s [end]
 d -> o -> g [end]

Dry Run Walkthrough

Input: words: ["cat", "car", "cars", "dog"], prefix: "ca"

Goal: Find all words starting with prefix "ca"

Step 1: Insert word "cat" into trie

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

Why: Add each letter as a node, mark end of word at 't'

Step 2: Insert word "car" into trie

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

Why: Share prefix 'ca', add 'r' node and mark end

Step 3: Insert word "cars" into trie

root -> c -> a -> t [end]
               ↓
               r -> s [end]

Why: Extend from 'r' node with 's' and mark end

Step 4: Insert word "dog" into trie

root -> c -> a -> t [end]
               ↓
               r -> s [end]
 d -> o -> g [end]

Why: Add new branch for 'dog'

Step 5: Search prefix "ca" in trie

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

Why: Traverse prefix letters to reach node where words start

Step 6: Collect all words starting from node 'a'

Found words: "cat", "car", "cars"

Why: Gather all words by exploring all branches from prefix node

Result:

root -> c -> a -> t [end]
               ↓
               r -> s [end]
 d -> o -> g [end]
Autocomplete results for "ca": ["cat", "car", "cars"]

Annotated Code

DSA Go

package main

import "fmt"

// TrieNode represents each node in the trie
type TrieNode struct {
	children map[rune]*TrieNode
	isEnd    bool
}

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

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

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

// Insert adds a word to 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 finds the node matching prefix
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 from 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
	}
}

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

func main() {
	trie := NewTrie()
	words := []string{"cat", "car", "cars", "dog"}
	for _, w := range words {
		trie.Insert(w) // insert each word
	}
	prefix := "ca"
	results := trie.Autocomplete(prefix) // get autocomplete
	fmt.Printf("Autocomplete results for \"%s\": %v\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 node reached

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

recursively explore all child nodes to find words

Output

Autocomplete results for "ca": [cat car cars]

Complexity Analysis

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

Space: O(n) where n is total letters in all inserted words for trie storage

vs Alternative: Compared to scanning all words for prefix (O(n * m)), trie is faster for large word sets by sharing prefixes

Edge Cases

Empty prefix string

Returns all words in trie since prefix is empty

DSA Go

if node == nil {
	return results // no words found
}

Prefix not in trie

Returns empty list because no words match

DSA Go

if node == nil {
	return results // no words found
}

Words with shared prefixes

Trie shares nodes for common prefixes, saving space and time

DSA Go

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

When to Use This Pattern

When you need to quickly find all words starting with a prefix, use a trie because it organizes words by shared prefixes for fast lookup.

Common Mistakes

Mistake: Not marking the end of a word, causing incomplete matches
Fix: Always set isEnd = true at the last node of inserted word

Mistake: Not checking if prefix node exists before collecting words
Fix: Return empty list if searchNode returns nil for prefix

Summary

Stores words by shared prefixes 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 to follow prefix nodes then collect all descendant words efficiently.