Asset 31Asset 32
Navigate back to the homepage

Prefix Trees in Go

Sebastian Ojeda
September 25th, 2019 · 2 min read

A trie is a data structure whose nodes store the letters of a given alphabet, and by which traversal of the nodes results in the retrieval of strings. If that’s a little difficult to follow, basically, a trie (pronounced try) is a search tree for words and strings. Typically, a trie is used because it provides fast search results with the ability to also match and search string prefixes. The time and space complexity for retrieval and insertion are both O(k), where k is the length of the string being inserted.

Let’s explore how to implement a trie in Go. In this example, we are building a dictionary that will allow us to put, get, and delete words as needed.

Start at the root

First, we need to define the structure of our tree and its nodes.

1type Trie struct {
2 root *Node
3 wordCount int
4}
5
6type Node struct {
7 prefix rune
8 parent *Node
9 children map[rune]*Node
10 isWord bool
11 results int
12}

We are storing the word count of the trie as well as the number of results (or words) that each node contains downstream. We also chose to use a rune for the simple face that:

“In general, a character may be represented by a number of different sequences of code points, and therefore different sequences of UTF-8 bytes.” -Go Blog

If you’re interested in learning more about the rune, read more about it on the Go Blog.

Now, if we want to initialize a new trie, we’ll need to include a trie and its corresponding root element, so let’s make a method for that.

1// New returns an initialized trie with a root node.
2func New() *Trie {
3 return &Trie{
4 root: &Node{
5 prefix: 0,
6 parent: nil,
7 children: make(map[rune]*Node),
8 isWord: false,
9 results: 0,
10 },
11 wordCount: 0,
12 }
13}

We can now create our soon to be dictionary, but how do we add a word? For this, we’ll create a method that allows us to loop over a word one character at a time and search our trie for that charcter, adding a new child node as necessary.

1// Put will add a new word to our trie, adding
2// new nodes as needed.
3func (t *Trie) Put(word string) {
4 node := t.root
5
6 for _, c := range word {
7
8 if n, ok := node.children[c]; ok {
9 node = n
10 n.results++
11 } else {
12 node.children[c] = &Node{
13 prefix: c,
14 parent: node,
15 children: make(map[rune]*Node),
16 isWord: false,
17 results: 1,
18 }
19 node = node.children[c]
20 }
21 }
22
23 node.isWord = true
24 t.wordCount++
25}

Now, for our dictionary to be of any use, we need to make sure that we can retrieve a string and then check if it is a valid word in our dictionary.

1// Get will try to return the node for the last
2// character in the string if found, as well as
3// a bool indicating whether or not it is a word
4// in the dictionary.
5func (t *Trie) Get(word string) (*Node, bool) {
6 node := t.root
7
8 for _, c := range word {
9
10 if n, ok := node.children[c]; ok {
11 node = n
12 } else {
13 return nil, false
14 }
15 }
16
17 return node, node.isWord
18}

Finally, once we have added words, we also need the ability to delete words from out trie. We must pay special attention to this method as we do not want to remove any nodes with more than on child. We must also remember to uodate the isWord value for any word-ending node whose word is being deleted.

1// Delete removes a word from a trie.
2func (t *Trie) Delete(word string) bool {
3 node, isWord := t.Get(word)
4
5 if !isWord {
6 return false
7 }
8
9 node.isWord = false
10
11 for node.parent != nil {
12 node.results--
13 if node.results == 0 {
14 delete(node.parent.children, node.prefix)
15 }
16 node = node.parent
17 }
18
19 t.wordCount--
20 return true
21}

Build the dictionary

That’s all there is to it. Pretty straight forward and not too difficult to implement in Go! Let’s see how our dictionary works.

1func main() {
2 dictionary := trie.New()
3
4 dictionary.Put("man")
5 dictionary.Put("manatee")
6 dictionary.Put("many")
7 dictionary.Put("canary")
8 dictionary.Put("mars")
9
10 fmt.Println(dictionary.wordCount) // 5 words
11
12 node, isWord := dictionary.Get("man")
13 if isWord {
14 fmt.Println(node.results) // 3 results
15 }
16
17 dictionary.Delete("manatee") // returns true
18
19 node, isWord = dictionary.Get("manatee") // returns nil, false
20}

So, there are a few features missing and improvements that could be made. As is, our implementation does not provide any additional benefits over other solutions such as a simple hash table. For this to truly be powerful, we’ll need to build onto this solution. For example, we could add a function to find all words that begin with a certain prefix, or fuzzy search functionality to more easily find what words. This is just a simple implementation after all, but definitely worth having in your toolbox if you ever come across a search tree problem.

More articles from Fullstack Go

Building an API endpoint with AWS Lambda and Go

Cloud computing has taken over. In this article, let's explore how to create an API endpoint with Go on AWS.

September 19th, 2019 · 3 min read

Handling JSON in Go

Go was built for the modern world. So let's talk about encoding and decoding the most popular data-interchange format.

September 18th, 2019 · 1 min read
© 2019 Fullstack Go
Link to $https://twitter.com/fullstack_goLink to $https://github.com/fullstackgo