← Back to DevBytes

Solving Implement Trie (Prefix Tree) in Go: Step-by-Step Guide

Introduction to the Trie Data Structure

A Trie, pronounced "try" and also known as a prefix tree, is a specialized tree-based data structure designed for efficient retrieval of keys in a dataset of strings. Unlike binary search trees that compare entire keys, a Trie breaks down each key into a sequence of characters, storing each character as a node along a path from the root to a leaf. This unique structure makes Tries exceptionally powerful for operations involving prefixes, autocomplete systems, and dictionary implementations.

In this tutorial, we will walk through implementing a complete Trie in Go, covering the core operations: insertion, search, and prefix matching. We will also explore best practices and performance considerations to help you write production-ready code.

Why Tries Matter

Understanding when and why to use a Trie is crucial for any developer working with string-heavy applications. Here are the key reasons Tries matter:

Compared to hash tables, Tries avoid hash collisions entirely and support prefix-based queries that hash tables simply cannot perform. Compared to binary search trees, Tries offer predictable O(L) performance for string operations rather than O(L log N).

Understanding the Trie Structure

Node Design

Each node in a Trie contains two essential pieces of information: a collection of child nodes (one per possible character) and a flag indicating whether the node marks the end of a complete word. In Go, we typically represent the children as a map for flexibility or a fixed-size array for performance when the alphabet size is known and small.

Operations Overview

A standard Trie supports three primary operations:

Implementing the Trie in Go

Step 1: Defining the TrieNode and Trie Structures

We begin by defining our node structure. Each node holds a map of child characters and a boolean flag indicating whether it represents the end of a word. We also define the Trie struct itself, which contains a pointer to the root node.

package main

import "fmt"

// TrieNode represents a single node in the Trie
type TrieNode struct {
    children map[rune]*TrieNode
    isEnd    bool
}

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

// Trie represents the prefix tree structure
type Trie struct {
    root *TrieNode
}

// Constructor initializes a new Trie
func Constructor() Trie {
    return Trie{
        root: NewTrieNode(),
    }
}

Using a map for children allows us to support any Unicode character without pre-allocating a large array. This is particularly useful when dealing with multilingual text or when memory efficiency is a priority over raw speed.

Step 2: Implementing the Insert Operation

The insert operation walks through each character of the input word. For each character, it checks whether a child node exists. If not, it creates one. Once all characters are processed, it marks the final node as the end of a word.

// Insert adds a word into the Trie
func (t *Trie) Insert(word string) {
    node := t.root
    for _, ch := range word {
        if _, exists := node.children[ch]; !exists {
            node.children[ch] = NewTrieNode()
        }
        node = node.children[ch]
    }
    node.isEnd = true
}

Notice that we iterate over the word using a range loop with rune type. This ensures correct handling of multi-byte UTF-8 characters, which is important for internationalized applications.

Step 3: Implementing the Search Operation

The search operation traverses the Trie following each character of the input word. If at any point a character is not found among the children, the word does not exist. If we reach the end of the word, we check the isEnd flag to confirm it represents a complete word rather than just a prefix.

// Search returns true if the word exists in the Trie
func (t *Trie) Search(word string) bool {
    node := t.root
    for _, ch := range word {
        if child, exists := node.children[ch]; exists {
            node = child
        } else {
            return false
        }
    }
    return node.isEnd
}

Step 4: Implementing the StartsWith Operation

The StartsWith operation is similar to Search, but it does not require the final node to be marked as the end of a word. It simply checks whether the prefix path exists in the Trie.

// StartsWith returns true if any word in the Trie starts with the given prefix
func (t *Trie) StartsWith(prefix string) bool {
    node := t.root
    for _, ch := range prefix {
        if child, exists := node.children[ch]; exists {
            node = child
        } else {
            return false
        }
    }
    return true
}

Step 5: Putting It All Together

Now let us combine all the pieces and write a main function to demonstrate the Trie in action.

package main

import "fmt"

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

func NewTrieNode() *TrieNode {
    return &TrieNode{
        children: make(map[rune]*TrieNode),
        isEnd:    false,
    }
}

type Trie struct {
    root *TrieNode
}

func Constructor() Trie {
    return Trie{
        root: NewTrieNode(),
    }
}

func (t *Trie) Insert(word string) {
    node := t.root
    for _, ch := range word {
        if _, exists := node.children[ch]; !exists {
            node.children[ch] = NewTrieNode()
        }
        node = node.children[ch]
    }
    node.isEnd = true
}

func (t *Trie) Search(word string) bool {
    node := t.root
    for _, ch := range word {
        if child, exists := node.children[ch]; exists {
            node = child
        } else {
            return false
        }
    }
    return node.isEnd
}

func (t *Trie) StartsWith(prefix string) bool {
    node := t.root
    for _, ch := range prefix {
        if child, exists := node.children[ch]; exists {
            node = child
        } else {
            return false
        }
    }
    return true
}

func main() {
    trie := Constructor()

    // Insert words
    trie.Insert("apple")
    trie.Insert("app")
    trie.Insert("application")
    trie.Insert("banana")
    trie.Insert("band")

    // Search for complete words
    fmt.Println("Search 'apple':", trie.Search("apple"))       // true
    fmt.Println("Search 'app':", trie.Search("app"))           // true
    fmt.Println("Search 'appl':", trie.Search("appl"))         // false
    fmt.Println("Search 'application':", trie.Search("application")) // true
    fmt.Println("Search 'banana':", trie.Search("banana"))     // true
    fmt.Println("Search 'bandana':", trie.Search("bandana"))   // false

    // Check prefixes
    fmt.Println("StartsWith 'app':", trie.StartsWith("app"))   // true
    fmt.Println("StartsWith 'ban':", trie.StartsWith("ban"))   // true
    fmt.Println("StartsWith 'cat':", trie.StartsWith("cat"))   // false
    fmt.Println("StartsWith 'bandana':", trie.StartsWith("bandana")) // true
}

When you run this program, you will see that searching for "apple" returns true because it was inserted as a complete word, while searching for "appl" returns false because it was never marked as a word ending. The StartsWith method correctly identifies that "bandana" is a valid prefix path even though it is not a complete word in the Trie.

Advanced: Collecting All Words with a Given Prefix

While the basic Trie implementation covers the standard LeetCode problem requirements, real-world applications often need to retrieve all words that share a common prefix. Let us extend our Trie with a method that collects all words starting with a given prefix.

// GetWordsWithPrefix returns all words in the Trie that start with the given prefix
func (t *Trie) GetWordsWithPrefix(prefix string) []string {
    node := t.root
    // Navigate to the node representing the end of the prefix
    for _, ch := range prefix {
        if child, exists := node.children[ch]; exists {
            node = child
        } else {
            return []string{} // prefix not found
        }
    }

    var results []string
    t.collectWords(node, prefix, &results)
    return results
}

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

You can test this method by adding the following to your main function:

fmt.Println("Words with prefix 'app':", trie.GetWordsWithPrefix("app"))
// Output: Words with prefix 'app': [app apple application]

Best Practices

Choose the Right Child Storage Strategy

Using a map for children is flexible and memory-efficient for sparse character sets. However, if you are working exclusively with lowercase English letters (a-z), a fixed-size array of 26 pointers can be significantly faster due to cache locality and the elimination of hash lookups.

// Array-based node for lowercase English letters only
type TrieNodeArray struct {
    children [26]*TrieNodeArray
    isEnd    bool
}

func (n *TrieNodeArray) getChild(ch byte) *TrieNodeArray {
    return n.children[ch-'a']
}

func (n *TrieNodeArray) setChild(ch byte, child *TrieNodeArray) {
    n.children[ch-'a'] = child
}

Handle Edge Cases Gracefully

Always consider what happens with empty strings. Inserting an empty string will simply mark the root node as a word ending. Searching for an empty string will return the isEnd flag of the root. Decide whether this behavior is appropriate for your use case and document it clearly.

Consider Memory Optimization

Tries can consume significant memory, especially with large datasets. For memory-constrained environments, consider a compressed Trie (also called a radix tree or Patricia trie) that collapses single-child chains into single nodes storing substrings rather than individual characters.

Use Pointers Consistently

In Go, always use pointers when working with Trie nodes. Using value types would cause copies of the entire subtree on every assignment, leading to both performance issues and incorrect behavior since modifications would not propagate back to the original structure.

Add a Delete Operation for Completeness

A production-ready Trie should support deletion. The delete operation must traverse to the target word's end node, unmark it, and then remove any nodes that are no longer part of another word.

// Delete removes a word from the Trie
func (t *Trie) Delete(word string) {
    t.deleteHelper(t.root, word, 0)
}

func (t *Trie) deleteHelper(node *TrieNode, word string, depth int) bool {
    if node == nil {
        return false
    }

    // Base case: reached the end of the word
    if depth == len(word) {
        if node.isEnd {
            node.isEnd = false
        }
        // Return true if this node has no children (can be deleted)
        return len(node.children) == 0
    }

    ch := rune(word[depth])
    child, exists := node.children[ch]
    if !exists {
        return false
    }

    shouldDeleteChild := t.deleteHelper(child, word, depth+1)

    if shouldDeleteChild {
        delete(node.children, ch)
        // Return true if current node is also deletable
        return len(node.children) == 0 && !node.isEnd
    }

    return false
}

Time and Space Complexity Analysis

Understanding the complexity of each operation helps you make informed decisions about when to use a Trie:

Conclusion

Implementing a Trie in Go is a rewarding exercise that deepens your understanding of tree-based data structures and string processing. The Trie's ability to perform prefix-based queries in time proportional to the query length, rather than the dataset size, makes it an invaluable tool for autocomplete systems, spell checkers, IP routers, and many other applications. By following the step-by-step implementation in this tutorial and adhering to the best practices outlined, you now have a solid foundation for building efficient, production-ready Trie solutions in Go. Whether you are solving the classic LeetCode problem or building a real-world search feature, the principles covered here will serve you well in your development journey.

— Ad —

Google AdSense will appear here after approval

← Back to all articles