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.