Trie

03 Jan 2015 in Go

Recently I’ve been working on a Trie implementation and I was a surprised and delighted from the simpicity of the data structure.

The first time I came across the Trie data structure, I was looking to solve a problem where a typeahead field had a really large dictionary to query and would chocke every time the search was triggered. After stumbling upon John Resig’s JavaScript implementation I wanted to see how it would translate in Go.

So without further ado, I’d like to share with what I found interresting about the Go version.

The Data Type

type Node map[rune]Node

Using self-referential type definition we are able to describe our tree nodes with nothing more than just built-in types.

Building an Index

Some parts are omited from the following example but this is where the magic happens.

func (node Node) Insert(s string) {
	for _, r := range s {
		if node[r] == nil {
			node[r] = New()
		}
		node = node[r]
	}
	node.End()
}

Notice how we walk down the tree for every rune in the string with node = node[r].

You’ve probably noticed the node.End() near the end. This inserts an invalid UTF8 character at the current node marking the end of a word.

Search

Searching is as simple as walking down the tree matching each letter in the search term to a key in each node. When we reach the end of the search term we have found a match, so the string indexed by the current node plus the strings from all its child nodes should be returned as results to the query.

func (node Node) Search(s string) (res []string) {
	for i, r := range s {
		if _, ok := node[r]; ok {
			node = node[r]
		}
		if i == len(s)-utf8.RuneLen(r) {
			return append(res, node.All(s)...)
		}
	}
	return res
}

The code is available on GitHub and the API docs at the usual place. I would love to hear your feedback. Enjoy!