Implement Trie (Prefix Tree)

Medium

The problem

Build a Trie, a structure that stores words. Insert(word) adds a word, Search(word) says whether exactly that word was added, and StartsWith(prefix) says whether any added word begins with prefix.

Example 1
Input: Insert("apple") Search("apple") Search("app") StartsWith("app") Insert("app") Search("app")
Output: true, false, true, true
Search("apple") is true. Search("app") is false because only "apple" was added. StartsWith("app") is true since "apple" begins with it. After Insert("app"), Search("app") is true. (Insert returns nothing, so it prints no result.)

Limits

1 ≤ word.length, prefix.length ≤ 2,000
Only lowercase English letters
Up to 30,000 calls in total

Write it in Go. Try for about 20 minutes on paper first, then open one hint at a time.

Try it here

Write Go. Common packages like fmt and sort are imported for you. Keep the function name and inputs the same.

Hints, one at a time

Nudge

Each node needs a way to reach the next letter and a flag saying "a word ends here".

The idea

Node = map/array of children + isEnd. insert walks/creates nodes; search needs isEnd; startsWith only needs the path to exist.

Target: O(len) per operation

Go function shape

type Trie struct{}
func Constructor() Trie
func (t *Trie) Insert(word string)
func (t *Trie) Search(word string) bool
func (t *Trie) StartsWith(prefix string) bool

Reference solution

Tested with go test. Try it yourself first, then compare.

// Trie stores words letter by letter, so words that share a beginning share a path.
type node struct {
	next map[rune]*node
	end  bool // a word finishes exactly here
}

type Trie struct{ root *node }

func NewTrie() *Trie { return &Trie{root: &node{next: map[rune]*node{}}} }

func (t *Trie) Insert(word string) {
	cur := t.root
	for _, r := range word {
		child, ok := cur.next[r]
		if !ok { // 1. no edge for this letter yet: create it
			child = &node{next: map[rune]*node{}}
			cur.next[r] = child
		}
		cur = child // 2. walk down one letter
	}
	cur.end = true // 3. mark the word's last node
}

// walk follows prefix from the root and returns the node it ends on, or nil.
func (t *Trie) walk(prefix string) *node {
	cur := t.root
	for _, r := range prefix {
		cur = cur.next[r]
		if cur == nil {
			return nil // 4. the path breaks: nothing stored starts with this prefix
		}
	}
	return cur
}

// Search reports whether the exact word was inserted.
func (t *Trie) Search(word string) bool {
	n := t.walk(word)
	return n != nil && n.end
}

// StartsWith reports whether any inserted word begins with prefix.
func (t *Trie) StartsWith(prefix string) bool { return t.walk(prefix) != nil }

Read the lesson: Tries