Implement Trie (Prefix Tree)

The problem

Build a class Trie that stores words and answers questions about them:

• insert(word) stores word.
• search(word) returns true if exactly word was stored before.
• startsWith(prefix) returns true if any stored word begins with prefix.

A trie, also called a prefix tree, is the usual way to do this. Each example below is a list of calls on one new Trie, with what each call returns (null for insert).

Example 1

Input: insert("cat"), insert("car"), search("cat"), search("ca"), startsWith("ca"), insert("ca"), search("ca")

Output: null, null, true, false, true, null, true

“ca” is not a word until it is inserted, but it is a prefix of “cat” and “car” from the start. The picture shows the trie at the end.

Example 2

Input: insert("dog"), search("do"), startsWith("dog")

Output: null, false, true

“do” was never inserted. A word counts as a prefix of itself, so startsWith(“dog”) is true.

Example 3

Input: startsWith("a"), insert("a"), search("a")

Output: false, null, true

An empty trie has no prefixes. After the insert, the one-letter word is found.

Constraints

• 1 ≤ word.length, prefix.length ≤ 2,000
• Words and prefixes use only the lowercase letters a to z.
• At most 30,000 calls in total.

Clarifying questions

Which letters can appear?

Only a to z. So each node needs at most 26 children, and an array of 26 works.

Is a word a prefix of itself?

Yes. startsWith("dog") is true after insert("dog").

Can the same word be inserted twice?

Yes. It changes nothing the second time.

Is there a delete?

No. Only insert and the two questions.

Spot the pattern

The questions are about starts of words. Words that start the same way can share the work of storing that start. A trie does exactly that: a tree where each edge is one letter, and each path from the root spells a prefix. Answering a question is a walk down one path, a small tree traversal.

Solutions

Brute forceO(N × L)TrieO(L)

Show all

This problem has two approaches. A hash set speeds up search but not startsWith, so it is no real middle step.

Brute force O(N × L)

The idea

Keep every inserted word in a plain list. To search, compare the word with each stored word. To answer startsWith, check whether any stored word begins with the prefix.

Inserting is cheap. But each question may look at every stored word, and each comparison may look at every letter. The list knows nothing about which words share a start.

Walkthrough

words

answers

empty

1. words = []
2. def insert(word): words.append(word)
3. def search(word):
4. for stored_word in words:
5. if stored_word == word: return True
6. return False
7. def startsWith(prefix):
8. for stored_word in words:
9. if stored_word.startswith(prefix): return True
10. return False

The brute force keeps every word in a plain list. Run the operations of Example 1 one by one.

← BackStep →Play1 / 9

Code

JSolution.javaPysolution.pyC#Solution.csGosolution.goTSsolution.ts

JavaPythonC#GoTypeScript

JavaPythonC#GoTypeScript

// Brute force: keep every word in a list and scan it on each question.
// With N words of length up to L, a search costs about N * L steps.
class Trie {
    private final List<String> words = new ArrayList<>();

    public void insert(String word) {
        words.add(word);
    }

    public boolean search(String word) {
        // Nothing is sorted or grouped, so every word must be compared.
        for (String storedWord : words) if (storedWord.equals(word)) return true;
        return false;
    }

    public boolean startsWith(String prefix) {
        for (String storedWord : words) if (storedWord.startsWith(prefix)) return true;
        return false;
    }
}

Big O

Time O(N × L)Space O(N × L)Here is how we get these from the code, in four steps.

1. 1 Find the work that repeats

   The lit lines compare one stored word with the question. They sit in a loop over the list. One compare reads up to L letters, where L is the length of the longest word, so it costs up to L steps. Let N be the number of stored words.

   1. words = []
   2. def insert(word): words.append(word)
   3. def search(word):
   4. for stored_word in words:
   5. if stored_word == word: return True
   6. return False
   7. def startsWith(prefix):
   8. for stored_word in words:
   9. if stored_word.startswith(prefix): return True
   10. return False

2. 2 Count the runs on Example 1

   The table counts the word compares on Example 1. The questions that find nothing, or find the word last, compare the most. search("ca") at the end compares all 3 stored words.

   1 + 2 + 1 + 3 = 7 word compares.

   operation	insert("cat")	insert("car")	search("cat")	search("ca")	startsWith("ca")	insert("ca")	search("ca")
   words in list	0	1	2	2	2	2	3
   words compared	-	-	1	2	1	-	3

   total

   1 + 2 + 1 + 3 = 7 word compares

3. 3 Write the count for any n

   A question that finds nothing compares all N words, and each compare can read up to L letters. So one question costs up to:

   N × L letter checks

   Check it with the last call of Example 1: N = 3, L = 3, so the bound is 9. The real compares read “cat”, “car” and “ca” for 2 + 2 + 2 = 6 letters, because each stops at the shorter word.

4. 4 Keep only the part that grows fastest

   There is nothing to drop: N × L has no smaller terms. Inserting costs 1 step, which is tiny next to a question. So one question is O(N × L), and it gets slower with every word inserted.

   So the time is O(N × L).

5. S Space: count the extra memory the same way

   The list holds every word, so it holds up to N × L letters. Nothing else grows. The space is O(N × L).

   words (all letters of all words)

Trie O(L)

The idea

Store the words in a trie. Each node has up to 26 children, one per letter, and a flag that says “a word ends here”. The root stands for the empty start.

• insert(word): walk down from the root, one letter at a time. When an edge is missing, add a new node. At the last letter, set the flag.
• find(prefix): walk down the same way. If an edge is missing, no stored word starts like this, so return None. Otherwise return the node at the end of the prefix.
• search(word): the walk must succeed and end on a flagged node.
• startsWith(prefix): the walk only has to succeed.

Words with the same start share one path. So the walk depends only on the length of the word, not on how many words are stored.

Why it is correct: a path from the root exists exactly when some inserted word starts with those letters, and its last node is flagged exactly when one inserted word ends there.

Walkthrough

answers

empty

1. def insert(word):
2. node = root
3. for letter in word:
4. if letter not in node.children:
5. node.children[letter] = TrieNode()
6. node = node.children[letter]
7. node.is_word = True
8. def find(prefix):
9. node = root
10. for letter in prefix:
11. if letter not in node.children: return None
12. node = node.children[letter]
13. return node
14. def search(word): return find(word) is not None and find(word).is_word
15. def startsWith(prefix): return find(prefix) is not None

A trie starts as one empty root. Each edge below it will stand for one letter, and shared starts will share one path.

← BackStep →Play1 / 12

Code

JSolution.javaPysolution.pyC#Solution.csGosolution.goTSsolution.ts

JavaPythonC#GoTypeScript

JavaPythonC#GoTypeScript

// Trie: words that share a start share one path of nodes, one letter
// per edge. Every operation walks one path, so L steps for length L.
class Trie {
    private static class Node {
        Node[] children = new Node[26];
        // Marks the end of a whole word, not just a prefix.
        boolean isWord;
    }

    private final Node root = new Node();

    public void insert(String word) {
        Node node = root;
        for (char letter : word.toCharArray()) {
            // Reuse the shared start. Add nodes only where it ends.
            if (node.children[letter - 'a'] == null) node.children[letter - 'a'] = new Node();
            node = node.children[letter - 'a'];
        }
        node.isWord = true;
    }

    public boolean search(String word) {
        Node node = find(word);
        // Reaching the end is not enough: a word must end exactly here.
        return node != null && node.isWord;
    }

    public boolean startsWith(String prefix) {
        return find(prefix) != null;
    }

    private Node find(String prefix) {
        // Follow one letter at a time. A missing edge means no stored
        // word starts this way.
        Node node = root;
        for (char letter : prefix.toCharArray()) {
            node = node.children[letter - 'a'];
            if (node == null) return null;
        }
        return node;
    }
}

Big O

Time O(L)Space O(N × L)Here is how we get these from the code, in four steps.

1. 1 Find the work that repeats

   The lit lines take one step down the trie: one lookup, maybe one new node, one move. That is a fixed job, so each letter counts as 1 step. Each operation runs this once per letter of its word.

   1. def insert(word):
   2. node = root
   3. for letter in word:
   4. if letter not in node.children:
   5. node.children[letter] = TrieNode()
   6. node = node.children[letter]
   7. node.is_word = True
   8. def find(prefix):
   9. node = root
   10. for letter in prefix:
   11. if letter not in node.children: return None
   12. node = node.children[letter]
   13. return node
   14. def search(word): return find(word) is not None and find(word).is_word
   15. def startsWith(prefix): return find(prefix) is not None

2. 2 Count the runs on Example 1

   The table lists the letters each call walks on Example 1. Every call walks exactly the letters of its own word: 3 for “cat” and “car”, 2 for “ca”.

   3 + 3 + 3 + 2 + 2 + 2 + 2 = 17 steps.

   operation	insert("cat")	insert("car")	search("cat")	search("ca")	startsWith("ca")	insert("ca")	search("ca")	total
   letters walked	3	3	3	2	2	2	2	17

3. 3 Write the count for any n

   An operation on a word of length L walks at most L edges, one per letter. A missing edge only makes find stop sooner:

   at most L steps per operation

   Check it with Example 1: each call took as many steps as its word has letters.

4. 4 Keep only the part that grows fastest

   There is nothing to drop: the count is already L. It does not depend on N, the number of words stored. A trie with a million words answers search("cat") in 3 steps.

   So the time is O(L).

5. S Space: count the extra memory the same way

   Each insert adds at most one node per letter, and only where no shared path exists yet. So the trie holds at most N × L nodes, and fewer when words share starts. Example 1 stores 3 words with 8 letters in 4 nodes plus the root. The space is O(N × L) in the worst case.

   one node per new letternode

Common mistake

Forgetting the end-of-word flag

Without the flag, search("ca") sees that the path c, a exists and returns true, even though only “cat” and “car” were inserted. The path says “a word starts like this”. Only the flag says “a word ends here”.

Tested at the limits

2,000 random call sequences of 1 to 15 calls, on words of 1 to 4 letters made of a and b, so prefixes, repeats and near misses come up all the time. The brute force and the trie gave the same answers every time.

Largest inputs	Result
30,000 calls (the most allowed) on 2,000-letter words that share their first 1,990 letters: 10,000 inserts, then 20,000 searches and prefix checks. Trie.	2.3 s for about 60 million letter steps, and about 20 MB of extra memory: about 100,000 trie nodes, which is the O(N × L) space.
2,000 calls of the same kind. Brute force.	0.04 s. At 30,000 calls it would compare each of 20,000 questions with up to 10,000 stored words.