Insert, search and prefix

Why it matters

Insert, search and prefix are the three operations behind almost every trie problem, and "implement a trie" is one of the most common interview questions for this data structure. Once you can write these methods from memory, the harder problems become small changes to code you already understand. This lesson builds the complete class step by step, runs it on real commands, and then adds the fifth operation, remove, which most beginners leave out.

Start with one helper: walking

Imagine giving directions in a large building. "Take the corridor marked c, then the one marked a, then the one marked r." Each turn follows one letter. If a corridor is missing at any turn, the place you want does not exist. If all the turns work, you have reached a spot. The question that matters is whether that spot is a real destination, and that is a separate question.

Every operation in this lesson begins with the same walk. You start at the root and, for each letter, move to that letter's child. If the child is missing, the walk stops and reports that the path does not exist. If every letter has a child, the walk ends on the node for the whole string. The walk returns that node, or None when a letter is missing. None is Python's value for "nothing here".

Writing the walk once, as a method, saves you from repeating the loop in every operation. Each operation then only decides what the final node means. A method name that starts with an underscore, such as _walk, is a convention that says "this is for use inside the class". Python does not enforce it.

Insert

Inserting a word means making sure every node on its path exists, then marking the last node as a word end. Start at the root. For each character, check whether the child for that letter exists. If it does not, create it. Then move into the child. After the last character, set the end-of-word flag.

Here is the insert of car, cat and care into an empty trie, step by step. Each row shows the step, the child that was created or reused, and where the cursor is afterwards. The cursor is the variable that points at the node you are on.

Ten steps, but only five nodes were created: c, ca, car, cat and care. Steps 1, 2, 3, 6 and 10 each created one of them. Steps 4, 5, 7, 8 and 9 reused nodes that already existed. That is the whole point of a trie: the second and third words cost little because they share the beginning of the first word.

Notice the order of the two actions in the loop: create the child if it is missing, then move into it. If you move first and create second, you move into None and Python raises an error. If you forget to move at all, the next letter is attached to the wrong node. The mistakes section below covers the first bug, and the first lesson covers the second.

Search and starts_with

Both methods walk down the tree one character at a time. If a child is missing, both return failure. The difference is what they require at the end.

• search(word): after the last character, the node must have its end-of-word flag set. The path is not enough.
• starts_with(prefix): reaching the node is enough, whatever its flag says.

A simple way to remember this is that starts_with asks whether a path exists, and search asks whether a word ends at the end of that path. Many beginner bugs come from answering the first question when the second was asked. The worked example below shows a program that does exactly that.

Counting words under a prefix

The fourth question is "how many stored words start with this prefix?" The easy way is to walk to the prefix node and count every word end below it. That is a full search of the subtree, which can be slow. The faster way is to keep a counter on every node: pass_count, the number of stored words that pass through it.

When you insert a new word, add one to the counter of every node on its path, including the root. The root's counter then equals the total number of words. The counter of the node for a prefix equals the number of words that start with that prefix. To answer count_with_prefix(p), walk to the node for p and return its counter. That is one walk, with no scanning of the subtree.

The counter must count distinct words. If you insert car twice, the nodes on its path must not be incremented twice. So insert first asks whether the word is already stored, and returns False without changing anything if it is. This is also why insert returns a boolean, which is Python's True or False value: the caller learns whether anything happened.

The complete class

Here is a complete program. It holds the commands in a string, one command per line, in the form command word. The commands are insert, search, prefix and count. The program prints the answer to each command. It uses two classes: Node holds the data for one node, and Trie holds the operations, so the structure is easy to test on its own. The match statement runs the case whose text equals the command, and case _ runs when none matches.

The program prints:

added
added
added
already there
True
False
True
False
2
3

Follow the commands one at a time. The first three inserts each print added, because none of those words was stored before. The fourth insert, car again, prints already there, because search("car") was already True, so insert returned False before it touched any counter. search car is True because the r node has its flag set. search ca is False because the walk reaches the a node, but that node's flag is off. prefix ca is True because the path to a exists. search dog is False because the root has no d child at all, so the walk returns None. count car prints 2, because car and care both pass through the r node, and its counter is 2. count c prints 3, because car, cat and care all pass through the c node.

Python prints the booleans as True and False, with a capital letter. Java prints them as true and false. The two languages differ only in this spelling.

Here is the same run as a table, showing what the store holds after each line. It is a good way to check your own trace before you run a program.

Lines 4 to 10 never change the store. Only the first three lines do. That is typical of a trie: most of its work is answering questions, and only some lines build the structure.

Read the code again with these outputs in mind. The Trie class holds all the logic. The loop at the bottom only splits each line and calls the right method. Keeping the two apart makes the class easy to test and reuse in the module assignment.

Worked example: why starts_with and search differ

The next program holds the lesson's most important trap. It stores two words and asks about a prefix of one of them. It runs two versions of the check side by side: a wrong one that only checks the path, and a right one that also checks the flag.

The program prints:

wrong: True
right: False

The words car and cat are stored, and then the program asks about ca. Here is the trace for the query ca. The walk starts at the root and moves to c, which exists because of car. It moves to a, which also exists. Then it stops, because there are no more letters. The walk returns the node for a, which is not None, so the wrong check says True. The right check then looks at that node's flag. The flag on a is False, because ca was never inserted as a word, so the right check says False.

The wrong check is not a little bit wrong. It answers a different question. It tells you whether anything in the dictionary starts with ca, which is what starts_with is for. A search that returns True for every prefix would mark ca as present even when the user has typed a fragment that is not a word, and a spell checker would stop flagging it.

This program uses a dataclass, which is a short way to write a class that mostly holds data. The @dataclass line above class Node tells Python to write the __init__ method for you from the fields listed below it. The field(default_factory=...) part makes a new list for each node. The lambda in it is a short unnamed function, here one that returns a new list. A plain default such as children: list = [None] * 26 is refused: dataclass raises ValueError: mutable default <class 'list'> for field children is not allowed: use default_factory. Python refuses it because one list would be shared by every node.

Worked example: a repeated word, with and without the check

The check at the top of insert exists for one reason, and this program shows it. It builds two tries from the same words. The naive one adds every word, even a repeat. The checked one skips a repeat. Both then count the words under car.

The program prints:

naive count for car: 3
checked count for car: 2

Trace the naive trie first. The first car adds one to the counters on c, a and r. The second car adds one more to each of them, so the r counter is 2. Then care adds one, so the r counter is 3. The count of 3 is wrong, because the trie holds only two distinct words under car, namely car and care.

The checked trie behaves differently on the second car. The function is_stored finds the flag already set, so add returns before any counter changes. The r counter stays at 1 after the first car, and then care makes it 2. The answer is 2, which matches the two distinct words.

The point is not that the naive version is slower. The naive version answers a different question. It counts how many times words were added, which is a count of inserts. The checked version counts how many different words are stored, which is the size of a set. Decide which question the problem asks before you choose the version.

Worked example: a table of queries

Start from car, cat and care, inserted in that order. The table asks the trie about several strings. The last column is what count_with_prefix returns for the same string.

Read the car row against the counts. The r node is on the path of car and of care, so its counter is 2. The care row has a counter of 1, because only care passes through the e node. The c and ca rows have counters of 3, because all three words start with c and ca.

The last three rows are the failures. In each one the walk stops at a missing child, and the method returns the failure value for its type: False for search and starts_with, and 0 for the count. A missing child is the normal case for most queries, so the walk must handle it cleanly. It does, because walk returns None the moment a child is missing and never looks further.

Notice that seven questions were answered from the same few nodes. That is the efficiency of the trie. Once the nodes exist, the same short walk answers every question about a prefix.

Why remove must update the counters

A shortcut is tempting: clear the flag and stop. search would then be correct, because it reads the flag. But the counts would be wrong. Take car and care, then remove car by clearing only its flag. search car is now False, which is right. count c still says 2, which is wrong, because only care is left. Nothing decreased the counter on c, so nothing can notice that a word has gone.

Clearing the flag also leaves dead nodes behind. In a trie where car was the only word, the nodes c, ca and car would stay in memory with no word below them. After many removals, a trie that only clears flags fills with paths that answer nothing. Decreasing the counters fixes both problems. The counters tell you which nodes have no words below them, and the cut removes them.

Remove

Removing a word is the reverse of inserting it, and it has one extra rule. You clear the end-of-word flag at the word's last node. Then you have to decrease the counter of every node on the path, because one fewer word passes through each of them. If a counter drops to zero, no stored word uses that node any more, so you can cut its link from its parent. Everything below it is also zero, so the whole branch can go at once, and the code stops there without updating the nodes below.

Here is a program that adds remove to the class. It runs a list of commands and prints the result of each. The lines starting with remove call the new method. The class is the same as before, with remove added at the end.

The program prints:

added
added
removed
False
True
True
removed
False
0
not stored

Trace the removals, because they are the hardest part. After the two inserts, the root counter, the c, a and r counters are each 2, and the e node for care has 1. When remove car runs, the root counter drops to 1, the c counter to 1, the a counter to 1 and the r counter to 1. None of them is zero, so the loop finishes at the r node and clears its flag. Now search car is False, and search care is still True. prefix car is True, because care still goes through the r node.

When remove care runs, the root counter drops to 0. The c node drops from 1 to 0, so the code cuts the link from the root to c. The whole chain is gone. prefix ca then finds no c child and returns False. count c has no node to read, so it returns 0. The last remove care finds nothing stored, so it prints not stored.

A second removal, step by step

Remove a shorter word from a trie that also holds a longer one. The trie holds car and care, so the r node has a counter of 2 and the e node has a counter of 1. Now run remove("care"), the longer word, to see what the cut does.

1. The method checks search("care"), which is True.
2. The root counter drops from 2 to 1.
3. The c counter drops from 2 to 1, then the a counter drops from 2 to 1, then the r counter drops from 2 to 1. None of them is zero, so the loop moves on.
4. The walk reaches the e node, whose counter was 1. It drops to 0, so the method sets the link from r to e to None and returns True.

The word car is untouched, because its flag sits on the r node and nothing changed that flag. The next commands show the result. count car prints 1, because only car is left under r. prefix care prints False, because the link to e is gone, and search care prints False too. count c prints 1. Those outputs all match the counters. Run the program with the commands insert car, insert care, remove care, count car, prefix care, search care and count c, one per line. The output is added, added, removed, 1, False, False and 1.

This trace shows the rule that matters most for remove. Removing a word never breaks a word that stays. It only lowers the counters on the path it used, and it cuts the nodes that no stored word uses any more.

What happens inside the computer

Each operation is a loop over the characters of one word. There is no recursion, so the call stack stays shallow. The stack is the record of which function called which, and each call uses a little memory. With a loop, the only memory that grows is the data in the trie itself.

The cost of a walk is one list index per character, so search, starts_with, count_with_prefix and the walk inside insert and remove all take O(L) time, where L is the length of the word. Insert and remove each walk the word more than once, once to check for duplicates and once to change the counters. That is still O(L), because two passes over a fixed string cost twice as much, which is a constant factor and not a growth rate.

Memory for the counters is one number per node. In Python a number is an object, not a fixed 4-byte value as a Java int is, so a counter takes somewhat more memory. It still adds a small amount to each node and does not change the growth rate of the structure.

What happens inside: the cursor

Look at the variable node in every method. In Python a variable is a name that points at an object. The line node = node.children[i] makes the name point at the child. It does not copy the child, and it does not make a new node. Afterwards node and the list slot both point at the same object. A walk is a sequence of these re-pointings, and the loop ends when node has moved as far as the word goes.

When remove cuts a link by setting node.children[i] to None, the node it used to point to loses one reference. Python frees memory by counting references. When a node's count reaches zero, its memory is released at once. The standard Python, called CPython, does this as soon as the last reference is gone. A node in this design never points back to its parent, so no cycle keeps it alive. That is why cutting one link is enough to free a whole branch: once the last reference to the top of the branch is gone, everything under it is released too. You never call a free function yourself.

Common mistakes

Returning the path check from search. If search is written as return self._walk(word) is not None, it becomes the prefix check. The wrong: True line in the worked example above shows exactly this bug. The fix is to read the flag at the end of the walk.

Moving into a child that does not exist. If the code does node = node.children[i] without first creating the child, node becomes None. The next letter then reads an attribute of None, and Python raises an AttributeError. The program below shows the related case, where a walk returns None and the caller uses the result without checking it. It reads is_word from the result of a walk for a word that is not stored:

class Node:
    def __init__(self):
        self.children = [None] * 26
        self.is_word = False
 
 
root = Node()
 
 
def walk(text):
    node = root
    for ch in text:
        node = node.children[ord(ch) - ord("a")]
        if node is None:
            return None
    return node
 
 
end = walk("dog")
print(end.is_word)

The output is a traceback. The file in this example is named main.py, so the path on the File line depends on where you save it:

Traceback (most recent call last):
  File "main.py", line 20, in <module>
    print(end.is_word)
          ^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'is_word'

Read the last line first. It names the type, NoneType, which is the type of None, and the attribute, is_word, that could not be read. The lines above it show where it happened, and the carets point at the part of the line that failed. The line number matches the line in your own file. The fix is to check for None first, as the search method does.

Here is the other case, the walk that moves into a missing child. The loop reads children from a node that is already None:

class Node:
    def __init__(self):
        self.children = [None] * 26
 
 
root = Node()
node = root
for ch in "xyz":
    node = node.children[ord(ch) - ord("a")]
print(node.children)

The output is:

Traceback (most recent call last):
  File "main.py", line 9, in <module>
    node = node.children[ord(ch) - ord("a")]
           ^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'children'

The root has no x child, so node becomes None after the first letter. The second letter then fails. The error appears one letter later than the real mistake, which is why the check for None belongs right after each step.

Counting duplicates. If insert skips the check for an existing word, then inserting car twice adds one to every counter on its path twice. The count car answer becomes 2 when it should be 1. Any counter in a set must count each member once, so the check at the top of insert is not optional.

Writing the wrong letter index. The index is ord(ch) - ord("a"), and the code must use the same index everywhere. Mixing ord(ch) - ord("a") in one method with ord(ch) in another gives an error or a missing word. The class above uses one helper, index_of, for this reason. Search the file for children[ and make sure each use goes through that helper.

Removing the flag and forgetting the counters. If remove clears is_word but does not decrement the counters, count_with_prefix keeps reporting the removed word. The output is quietly wrong, with no exception to warn you. Print the counter after each remove in a test to catch it.

Cutting a link too early. If remove cuts a link while another word still uses it, the other word disappears. The check against zero must come after decrementing, and it must use the counter of the node being cut, not the parent's.

Debugging walkthrough: a test table

When a trie gives a wrong answer, do not guess. Build a small table of commands and the answers you expect, run the program on it, and compare line by line. The first wrong line tells you which operation to examine.

Take the removal test from above. Suppose the output had shown prefix car as False on line 6. You would look at what should be true at that moment. care is still stored, so the r node must still exist, and the walk for car must reach it. A false answer means either the r node was cut too early or the walk stopped before the end. The dump from the first lesson would show which, because a missing node is missing from the printed tree.

Use this order when you debug. First check the walk, because a broken walk breaks every operation. Then check the flag, using search on a word you know is stored. Then check the counters, using count_with_prefix on a prefix you know. Each check narrows the problem to one part of the class.

A method that prints the counters

When a count looks wrong, print the counter on every node. This method walks the whole trie and prints each node's path and counter. It goes inside the Trie class, next to the other methods. It is written as a method that takes the node and the path so far.

def show(self, node, path):
    if node.pass_count > 0:
        print(path, node.pass_count)
    for i, child in enumerate(node.children):
        if child is not None:
            self.show(child, path + chr(ord("a") + i))

Call trie.show(trie.root, "") after inserting car, cat and care. The output is:

 3
c 3
ca 3
car 2
care 1
cat 1

The first line holds only the root's count. It starts with a space, because the path is an empty string, so it looks like a line with a space and a 3. The helper visits the children of each node in alphabetical order, because it loops over the list from slot 0 to slot 25. So car comes before its child care, and cat comes last. Compare each line with the table above. A counter that does not match the words below its node is the bug to look for.

The helper takes time proportional to the number of nodes, so keep it out of the final program. Use it while you test, then delete it.

Variations and edge cases

An empty word. insert("") sets the flag on the root. search("") then returns True, and count_with_prefix("") returns the root counter, which is the total number of words. Decide whether your problem allows empty words before you write the check.

A prefix longer than any word. Suppose only car is stored. starts_with("carrot") walks c, a and r, and then looks for a second r child under car. There is none, so the walk returns None and the answer is False. Check the last letter, not only the first few.

A word inside another word. If you store car and then care, the flag on r stays True. Removing care leaves car intact, because the counter on r goes from 2 to 1 and the flag is not touched.

Insert after remove. If you remove a from a trie that held only a, the root counter and the a counter both reach zero, and the a link is cut. Inserting a again creates a fresh node and sets its flag. search("a") then returns True. A correct removal leaves no trace behind.

Mixed case. The class assumes lowercase letters only. A capital letter gives a negative index. Python does not raise an error for a negative list index inside the list's range, because it counts from the end. So ord("T") - ord("a") is -13, which is slot 13, and the word is stored under the wrong letter with no warning. Convert input with word.lower() before calling any method, unless the problem says mixed case is part of the data.

Removing a word and adding it back. Remove car, then insert car again. Any node that the removal cut is rebuilt by the insert, and the counters count the word again. This case exercises the cut and the rebuild in one step, so test it.

Counts after many removals. A counter is only as good as the updates that change it. If one path forgets to decrement, the counts drift away from the real words, and the drift grows with each removal. A good test inserts a list of words, removes them all in a shuffled order, and checks that the root counter is zero at the end.

Interview-style questions and real uses

"Implement a trie" is usually asked with three requirements: insert, search and starts_with. A strong answer writes the walk helper first, then the three operations, each in a few lines. Mention that search needs the flag and starts_with does not. Then mention the complexity: O(L) for each operation, with L the length of the word.

A follow-up is often "how would you count words with a prefix quickly?" The answer is the per-node counter described above. The time per query is O(L), instead of scanning every word below the prefix.

Another common follow-up is "how would you delete a word?" The answer is the decrement-and-cut method, and the reason it is safe: a counter of zero means no stored word passes through that node.

Real systems use the same ideas. Search engines keep prefix structures for suggestions. Address books and phone contact searches find people by their first letters. Compilers use tries to match keywords while they read source code. In each case, the trie answers the prefix question quickly and keeps the answer correct when items are added or removed.

Explaining the code in one minute

An interviewer may ask you to walk through your code. A clear order helps the listener follow along. Start with the node: its children and its flag. Then the walk, the one loop that every operation shares. Then insert, which creates missing nodes and sets the flag. Then search and starts_with, which differ only in the last check. Then the counters, which make prefix counts fast. Finally remove, which lowers the counters and cuts a link when one reaches zero. Each part builds on the one before. If you run short of time, the walk and the flag are the two parts that matter most.

Practice

1. Trace search("care") on a trie holding car, cat and care. Say which nodes you visit and what the flag says at the end. Hint: the walk visits c, ca, car, care.

Show the answer: The walk visits c, then ca, then car, then care. The last node is care, and its flag is set because it was inserted. The answer is True.

2. After inserting car, cat and care, what does count_with_prefix("ca") return? Explain using the counters.

Show the answer: It returns 3. The a node lies on the path of all three words, so each of them added one to its counter.

3. Write a method is_empty(self) that returns True when no words are stored. Hint: the root counter is the total number of words.

Show the answer: Write return self.root.pass_count == 0. The root counter goes up on every insert and down on every remove, so it is zero exactly when the trie is empty.

4. A student removes the duplicate check from insert and inserts car twice. What does count_with_prefix("car") return, and what should it return?

Show the answer: It returns 2. It should return 1, because only one distinct word was stored. Each insert adds one to each counter on the path, so the second insert double-counts. The check if self.search(word): return False prevents this.

5. Write a method count_words(self) that returns the number of distinct words stored. Use the existing fields and no loops.

Show the answer: Write return self.root.pass_count. The root is on the path of every word, so its counter is the total.

6. Make search ignore case by converting the input to lowercase. Which line changes, and why must the change come before the walk?

Show the answer: Change the start of search to word = word.lower(), and do the same in insert. It must come before the walk, because the walk uses ord(ch) - ord("a") as the index, and a capital letter would give a negative index. Python would not raise an error for that, so the word would be stored under the wrong letter.

7. Trace the removal of car from a trie holding only car. Which counters change, and which link is cut?

Show the answer: The root counter goes from 1 to 0, and so does the c counter. The loop reaches c, sees its counter is zero, and cuts the root's link to c. The a and r nodes are no longer reachable, so their counters are left at 1. The remove method returns True.

8. Explain why insert calls search before it walks, rather than walking first and checking the flag at the end.

Show the answer: The walk in insert creates nodes as it goes, and it adds one to counters as it goes. If the word is already stored, the code must return without changing any counter, and the check has to happen before the first change. Calling search first is the simplest way to keep the counters correct.

9. A trie holds only ab. You call remove("ab"), then insert("a"). What does search("a") return?

Show the answer: It returns True. The removal cut the links down to nothing, so insert("a") creates a fresh a node, and sets its flag. The new word is stored normally.

10. Explain why starts_with does not need the flag check that search has.

Show the answer: starts_with only asks whether any stored word begins with the prefix. If the walk reaches a node, some stored word passes through it, because every node in the trie lies on the path to a stored word. The flag only says whether the node itself is a whole word, which is not what the question asks.

What you now know

• The walk follows one letter per step and returns the node for a string, or None if a letter is missing.
• Insert creates missing children, moves into them, and sets the end-of-word flag. It returns False for a word that is already stored.
• Search checks the path and the flag. starts_with checks the path only.
• A counter on each node gives the number of stored words under it, so prefix counts take one walk.
• Remove clears the flag, decrements the counters, and cuts a link when its counter reaches zero.
• Watch for the common bugs: a missing flag check, a double-counted duplicate, and a None used without a check.

In the next lesson you will use the same nodes to list words in alphabetical order, and to search a grid of letters for many words at once.