Implement Trie (LeetCode #208)

Image by Maria Kavelashvili on Unsplash

This article is part of a series from Nathan Thomas, a full stack software engineer working in San Francisco, CA. Other recent articles from him include Building Your Own Bitcoin Node and Jump Game.

Introduction

If you’re looking for a quick walkthrough of an optimal solution to the LeetCode Implement Trie problem, you’re in the right spot.

This is a question in the Blind 75 LeetCode code challenge list, a group of questions put together by a tech lead at Facebook that’s been touted as a great way to prep for interviews.

Don’t beat yourself up too much if you’re here because it gave you a hard time. Everyone has a hard time with coding challenges sometimes.

Let’s get it started!

Understanding the Problem 🕵🏻‍♂️

The first thing we should always try to do is understand all the rules for our challenge before writing a single line of code (see Polya’s Problem Solving Techniques):

1. We’re supposed to implement a Trie (which is pronounced as “try”) or “prefix tree” (if you want to read up on what it is, read the introduction for this Wikipedia link)
2. We should build the following functionality for it:

• Trie() instantiates the trie object
• insert(word: str) inserts the word into the Trie
• search(word: str) returns True if the word exists or False otherwise
• startsWith(prefix: str) returns True if there is a previously-inserted word in the Trie which the prefix prefix or False otherwise

I’ll be using Python for the answer for this question since it’s a bit like a universal language — whether you prefer JavaScript, Java, Golang, C, or something else, everyone can “kind of” read Python.

We’ll also be jumping straight into an “optimal” solution. If you’re anything like me, you already attempted this problem before searching for an answer. I’m going to make sure you get a good one.

Okay. I think that’s enough housekeeping items for us to get started!

Initial Problem Setup 🏗

If you’ve seen a Tree data structure before, you’re just a hop/skip/jump away from learning what a Trie is (and if you haven’t, read the introduction on this Wikipedia page).

It turns out that a Trie is very, very similar to a Tree. The key difference is that a Trie can have many children for each node whereas the Tree can only have two (a left and a right).

In a Trie, each node may contain many children that each represent another character in a chain that you are building. In the Trie below that I’ve written for you, you can see that there are multiple words contained in it if you start from Root and then progressively move down the nodes one-by-one:

In the Trie above, we can search through the nodes progressively to check if a word is contained in it.

For instance, searching the Trie above for “TREE” would return a true boolean:

Because of this, Tries are a powerful way to solve problems that involve words and word prefixes.

For our problem, the first things we need to do for our initial setup are to:

1. Define a Node class for each of the nodes (like in the diagram above)
2. Build the actual Trie's constructor to handle instantiating a new Trie

To start off with, we’ll define the Node class so that it has a constructor with a value, children dictionary, and word_totals value (more on that later):

Next, we’ll create a class called Trie that has a constructor. In that constructor, we’ll merely instantiate a root node that has a value of None:

This is typical for a Trie; we never have a value from the root node, and the children of the root node are the first letter in any words that might exist in the Trie.

From here, we’re ready to fill out our three methods!

Let’s do it. 💪🏻

Building the Insert Method

The insert method will look very familiar to you if you’ve ever had to build out methods to parse through a Binary Search Tree (BST). The big difference here is that we don’t have left or right values. We’re looking in a children dictionary each time we arrive at a new node.

First, let’s create a variable called current and initialize it with self.root from the Trie:

We’ve also set up a variable called word_index that allows us to independently track what the current index of the word is that we’re looking for.

Next, we’ll build out a while loop that runs for as long as word_index is less than the length of word. Inside this, we’ll check for every node we visit if the word[word_index] current exists in that node’s children and then create it as a new node if it doesn’t:

This means that we’ll get to reuse a node if it exists or create it otherwise.

Nice 🎉

Next, we should update current to be this either pre-existing or newly-created node:

Notice that we’re also incrementing word_index here as well.

Finally, the last thing we did up there at the end of our code was to increment the word_totals for the last node; this is to recognize that we’ve added one word to the Trie that ends in this node.

This will be useful for us very soon.

Our insert method runs in O(n) time (where n is the number of characters in word) and has O(1) space complexity (due to only needing a current and word_index pair of variables).

Building the Search Method

If you read through every step of the insert method above, the search one is going to look largely the same to you.

First, we do our variable setup and construct the start of our while loop:

Notice that we’re merely returning False if we hit a case where a word[word_index] character value does not exist in the current.children dictionary.

Next, we should do the same updates to current on every loop as well as incrementing word_index:

The only big thing here that differs from what we saw in the insert method is the return statement at the end.

We’re checking if the word_totals (which we incremented at the end of the insert statement if a word had finished being inserted on that node) is greater than 0 (which means there is a word present there).

Our search method runs in O(n) time (where n is the number of characters in word) and has O(1) space complexity (due to only needing a current and word_index pair of variables).

Cool!

That’s two methods down. Ready for the last one? 🚀

Building the SearchWith Method

This one is where the real money lies with this coding challenge. We’re going to have to pull out the old Depth First Search to make it happen.

We’ll start off our startsWith method in the exact same way as the search method. In fact, I’ll just plop the code down since it’s 100% the same:

Look familiar? 🤔

However, this is where the similarities cease. We can only use this part of our Trie traversal up until the end of the prefix. The whole point is that we should find if any words exist that start with this prefix.

This means that we have to roll up our sleeves and start searching through the remaining nodes from where the prefix left off.

First, let’s go ahead and create a stack variable and initialize it with a list that (for now) just contains whatever node current currently is:

We’ll be adding values to the front of the stack here and popping them off the front, meaning that we’ll be digging down into our Trie in a depth-first search pattern (technically in a somewhat “post-order” depth-first search manner).

We’ll next add another while loop that will run for as long as stack has a length that’s greater than 0. For each iteration, we’ll pop off the top of the stack and assign it to our previous current variable:

If the current node popped off the stack has a word_totals of greater than 0, we know we’ve found a word. In this case, we can just return True.

If not, we should go ahead and add all the children in current.children to the top of the stack:

If we ever exit out of our while loop without finding any word_totals on any of our nodes, we can merely have a base statement that returns False.

Boom. There we go! 🙌🏻

This method will run in a worst-case time complexity of O(V) where V is the number of nodes in the Trie. The space complexity is potentially O(V) as well due to the stack that we’re using.

Conclusion

Nice work. This one teaches you a fundamental data structure, so hopefully going through this has cemented it in your brain!

Keep an eye out for the rest of these Blind 75 coding challenges articles. It’s going to be a fun ride.

Feel free to reach out (using the links below) if you want to talk about coding, cool tech, or anything else (I really like book recommendations).

Thanks for reading. 🔥

Nathan

(GitHub, LinkedIn, Twitter, and Personal Site)