Autocomplete and word search on a grid in Python

Why it matters

Two famous trie problems use everything from the last lesson. Autocomplete asks for the words that start with what the user typed, in a useful order, and it must stay fast when the dictionary is large. Word Search II asks for every word from a long list that can be traced through a grid of letters. The first is a single walk followed by a walk over a subtree. The second is a depth-first search over the grid that uses the trie to cut off hopeless branches early. Both are common in interviews, and both teach a pattern you will reuse: do the work once in a shared structure instead of once per word.

Autocomplete: collect the words under a prefix

Suppose the store holds car, card, care, cat and dog, and the user has typed ca. The answer is every stored word whose path begins with c and a. In the trie, those words sit in the subtree below the node for ca. So the job has three parts.

1. Walk to the node for the prefix. If the walk fails, nothing starts with that prefix, and the answer is an empty list.
2. From that node, visit everything below it and collect each node whose end-of-word flag is set.
3. Build each word as you go. The word is the prefix plus the letters on the path from the prefix node down to the flagged node.

Step 2 is a depth-first search, usually shortened to DFS. A DFS goes as deep as it can down one branch, then backs up and tries the next branch. It is the natural way to walk a tree, because it follows one path to its end before it moves on. The alternative, a breadth-first search, visits all nodes at one depth before going deeper. For autocomplete a DFS is simpler to write, and the order of the children gives the output order you want.

The alphabetical order comes for free. If the code visits children from slot 0 (a) to slot 25 (z), the words come out in dictionary order. Within one node, the word that ends at that node comes before the words below it, because car is a prefix of card, and the dictionary order puts the shorter word first.

Here is the complete program. The words and the prefix are in the code, so the program needs no input. It prints the list of matches.

The program prints:

['car', 'card', 'care', 'cat']
[]

The first line is a Python list. Python shows each string in quotes, and a list of strings prints with square brackets and commas. The second call asks for x. No stored word starts with x, so the walk fails at the root and returns an empty list.

Follow the DFS by hand. It starts at the ca node, which is not a word, so nothing is added. It then visits the child r. That node is the word car, so car is added. The car node has two children, d and e, so the DFS goes to card first and adds it. Then it returns and goes to care, which is added. Back at ca, the next child is t, and the cat node is a word, so cat is added last. The order car, card, care, cat is the alphabetical order, and it came from the order in which the code tried the children.

The dog branch is never visited, because the walk for ca never goes near it. That is the point of the trie: the work depends on the size of the answer and the length of the prefix, not on the size of the dictionary.

Why depth first and not breadth first

A breadth-first search would also work for autocomplete. It visits the nodes one level at a time, using a queue, which is a line of waiting items served first in, first out. The trouble is the output order. Under ca, the first level holds car and cat, and the second level holds card and care. A breadth-first walk would print car, cat, card and care. That is level order, not dictionary order, and a user expects the dictionary order.

A depth-first walk prints car, card, care and cat, because it finishes everything under car before it moves on to cat. Depth-first also needs less memory on a wide tree. Its memory is the current path, which is at most as long as the longest word, while a breadth-first queue can hold a whole level of the tree. With a limit on the result, that matters, because the search may stop early and a breadth-first queue would have held far more nodes than were needed.

Building the word on the way down

The collect function keeps one list of letters, path, for the whole walk. When it goes down a child, it appends the letter with path.append. When it comes back up, it removes that letter with path.pop(). Each word is turned into a string with "".join(path) only when it is added to the output.

Java code for this step often uses a StringBuilder, which is a string you can change in place. Python strings cannot be changed in place. Writing prefix + ch makes a new string every time, which is correct but copies the text again on each step. A list of single letters, joined once per word, avoids that copying. The pattern is the same as the StringBuilder version: make a choice, explore it, and undo it before trying the next choice.

This is called backtracking. The same pattern appears in the grid search later in this lesson.

The recursion itself is worth a closer look. Each call to collect waits for its child calls to finish before it returns, so Python keeps a record of the waiting calls. That record is the call stack. The depth of the stack equals the length of the current path, so for words of up to 20 letters the stack is never deeper than about 21 calls. Python allows 1,000 nested calls by default, and going past that raises a RecursionError. A trie for ordinary words stays far below that limit.

What the cost really is

Reaching the prefix node takes O(L) steps, where L is the prefix length. After that, the DFS visits nodes below the prefix node. Here is a useful fact about any trie: every node exists only because some stored word passes through it, so every node leads to at least one word. The DFS therefore never walks into a branch that produces nothing. Each node it visits adds to the answer in some way, either by being a word or by leading to one.

With no limit, the DFS visits every node below the prefix, so the cost is O(L + size of the subtree). Real systems rarely want every match. They want the first few. That is the next section.

Capping the result

Suggestions in a search box show a handful of words, not thousands. A limit of k words lets you stop early. The change is small: before visiting a node, check whether the list already has k words. If it does, return at once. Every call above the stopping point returns quickly too, because each one checks the same condition before doing any work.

The next program adds a limit to suggest. It also counts how many trie nodes the search visits, so you can see the saving. The count is a module-level variable, visited. A function that changes a module-level variable must say so with the global keyword, or Python would create a new local variable with the same name instead.

The program uses math.inf, which is Python's value for infinity. A list is never longer than infinity, so passing math.inf as the limit means "no limit". Java's version uses Integer.MAX_VALUE for the same purpose.

The program prints:

['car', 'card', 'care', 'cat']
visited: 5
['car', 'card']
visited: 3

Count the visits by hand to see that the program does what you expect. With no limit, the DFS visits ca, car, card, care and cat, which is five nodes. With a limit of 2, it visits ca, car and card. After card the list holds two words, so the call for care returns before it counts a visit, and the call for cat does the same. The total is three.

The saving looks small here because the example is small. On a dictionary with a hundred thousand words, a prefix such as ca can have thousands of words below it. A limit of 10 then stops the search after a few hundred nodes, not after the whole subtree.

Even with a limit, the cost is bounded. Every node the DFS visits lies on the path to a word that ends up in the list, because the search only enters a branch that has a word below it and it stops as soon as the list is full. A word of length M has at most M nodes on its path below the prefix. So with a limit of k, the search visits at most about k * M nodes after the prefix walk. That is a useful number to state in an interview.

The calls, one by one

The visit counter in the last program is easier to trust when you see each call. The table lists every call to collect for the prefix ca with a limit of 2. The last column shows the list after the call's own work.

Call 1 is the ca node itself. It is not a word, so nothing is added, but it still counts as a visit. Call 2 goes to car, which is a word. Call 3 goes down the d child of car to card, which is also a word, and the list is full after it. Calls 4 and 5 are the e child of car and the t child of ca. Both return at the top, before they count as visits. The visit total is 3, which is the number the program prints.

Word Search II: many words in one grid

Now the harder problem. You are given a grid of letters and a list of words. You must find every word that can be traced through the grid. Two neighbouring cells are those that touch side to side or up and down. Diagonal steps do not count. A word may not use the same cell twice.

The slow method searches the whole grid once for each word. A grid of 16 cells and a list of 100 words means 100 separate searches. Each search tries every start cell and explores every neighbour path. Most of that work repeats across words: two words that start with oa share the same first steps through the grid.

The better method puts all the words into one trie and runs one DFS from each cell. As the DFS walks the grid, it also walks the trie. Each step on the grid must match a child in the trie. If the letter under the DFS has no child at the current trie node, no stored word can continue along this path, so the DFS stops right there. Cutting off a branch of a search that cannot succeed is called pruning. The trie tells you when a branch is hopeless, and that is what makes this method fast.

Here is the complete program. The grid is a list of rows, and each row is a list of one-letter strings. The program builds the trie from the word list, runs one search from each cell, and prints the words it finds in alphabetical order, one per line.

The program prints:

eat
oath

Why these two and not the others? oath can be traced from the top-left o to a on its right, then down to t, then down to h. eat can be traced from the e on the right edge, left to a, then left to t. pea needs a p, and the board has none. rain starts at the r in row 2, column 3, whose neighbours are e, v and k. None of them is a, so the trie has no child for them under r and the search gives up after one step.

Python's set plays the role that Java's HashSet plays. It stores each word once. The sorted call returns the words in alphabetical order, which is the order the output needs. A plain list would show duplicates if the same word were found twice, which is why the code also clears node.word.

The grid search is a depth-first search on a graph, where each cell is a node and each neighbour touch is an edge. Our search adds one extra check at each step: the trie child must exist. That check is the pruning.

Worked example: tracing eat through the grid

This trace follows one complete search, so you can see the marks and the trie walk together. The word is eat, and the board is the one from the example above.

The outer loop starts at row 0, column 0, and works through the cells in order. The o at the top left starts a search that finds oath, and it marks several cells on the way. The search for eat begins later, when the loop reaches row 1. Here is what happens at the two e cells in row 1.

1. The loop reaches row 1, column 0, which holds e. The root has a child for e, because eat is stored, so the search goes on. It marks this cell as used and tries the neighbours it has: o above, i below and t to the right. The e node has no child for o, i or t, so each of those calls returns at once. The search from this cell finds nothing, and the cell is unmarked on the way out.
2. The loop reaches row 1, column 3, which holds e again. The root's child for e is the same node as before, so the walk starts at the e node. This cell is marked.
3. The neighbour to the left is a, at row 1, column 2. The e node has a child for a, which is the node for ea. The search moves there and marks the a cell.
4. From the a cell, the neighbour to the left is t, at row 1, column 1. The ea node has a child for t, which is the node for eat. That node holds the word eat, so the search adds it to the result and clears the word field.
5. The search tries the rest of the neighbours of t, finds no further words, and returns. Each cell on the path is unmarked as the calls return.

The trace makes one point clear. The search for eat needed the t cell at row 1, column 1 to be free. That cell was on the path of the oath search earlier, and the restore step freed it. The next section shows what happens when the restore step is missing.

How the marking works

The line board[r][c] = "#" marks the cell as used on the current path. Without it, the search could step back and forth between two cells and form a word that reuses a cell. The line board[r][c] = ch after the loop undoes the mark, so other paths can use the cell again. The mark is a choice made before the recursive calls, and the undo happens after them. This is the same backtracking pattern as in the autocomplete code.

The check if ch == "#": return must come first. If you look up parent.children[ord(ch) - ord("a")] while the cell holds #, the index is ord("#") - ord("a"), which is 35 - 97, or -62. A list of 26 items has no slot -62, so Python raises an IndexError. The mistakes section shows that message.

Why the word is stored on its node

The trie node for each word holds the full word as a string, not a flag. That lets the search report a word as soon as it reaches the word's last cell, without rebuilding the string from the path. The code then sets the field back to None, so the same word is never reported twice if two paths find it.

Cost

For one start cell and one word of length L, a naive search can try up to 4 neighbours at each step, so it explores up to 4 to the power L paths. That number grows fast. The trie limits the search to paths that match a stored prefix, and most paths stop after one or two letters on a real board. The real saving comes from sharing: the same grid path that serves oath also serves oat and oa, and the trie explores it once, not once per word.

For the interview answer, say this. Build a trie from all the words. Run one DFS per cell. Follow the trie as you go, and stop when there is no child. Mark cells while they are on the path and unmark them after. The running time depends on the board and the trie, and it is far smaller than searching once per word.

Sharing a path, with an example

Take the two words oath and oat. Both start at the top-left o, then go to the a beside it, then to the t below that. A search for one word at a time walks those three cells twice, once for oath and once for oat. In the trie, the path for oat is the first three steps of the path for oath, so a single DFS walks those cells once. At the t node it sees that oat is a word and reports it, and then it keeps going down the h child to report oath. The shared steps are paid for once, and that saving grows with the number of words that share a beginning.

Explaining Word Search II in one minute

An interviewer may ask you to explain the method without writing code. A short order works well. Say what the input is: a board and a list of words. Say the simple method, one search per word, and why it repeats work. Then say the better method: put all the words in one trie, run one search from each cell, and follow the trie as you go. Name the pruning: stop when there is no matching child. Name the marking: mark a cell while it is on the current path, and unmark it when the recursion returns, so each path sees a clean board. Finish with the shared path: the trie explores each common beginning once, however many words use it.

Common mistakes

Forgetting to restore the cell. If board[r][c] = ch is missing, a cell used by one path stays # for every later path. Later words that need that cell are missed, and the output is missing words with no error. This is the most common Word Search bug. The next section shows the output.

Checking bounds after the recursive call. If the code calls search(nr, nc, node) before checking that nr and nc are inside the grid, the call reads past the board. The message is an IndexError, and it appears when the row or column is one past the last index. Python has a second trap that Java does not have. A negative index counts from the end of the list, so board[-1] is the last row and no error is raised. A search that steps up from row 0 would read the bottom row without any sign of a problem. Always check the bounds first. The check 0 <= nr < len(board) covers both cases.

The next snippet shows the silent wrap:

The output is:

b
index 2 is past the end

Index -1 gives the last row, b, and no error. Index 2 is past the end, and asking for it directly would raise IndexError.

Returning a word twice. A word can be reached by two different paths on the board. Without the step that sets node.word to None, the same word is added twice. A set hides this here, but a list would show duplicates.

Building the autocomplete string without removing the last letter. If collect appends a letter on the way down and forgets to remove it on the way up, the next sibling gets the wrong prefix. The output then has words with letters stuck to the front or back. The path.pop() call is the removal, and it must run after the recursive call returns.

Using the wrong limit check. Checking the size of the list after the loop, rather than before the node is visited, lets the search visit extra nodes. Put the check at the top of collect.

Reading the grid from input. The grid in this lesson is a list literal, not input. If your input has spaces between letters, read each row with split(). If your grid has no spaces, such as oaan, each row is one token and you must turn it into a list with list(row). Read the input format before you write the reader.

Debugging walkthrough: the word that is missing

Suppose you run Word Search II and oath is missing from the output. Work through these steps in order.

First, check the trie. Print the word list you inserted, and confirm that oath is there with its full spelling. A missing h in the list is a typing mistake, not a logic bug.

Second, check the start. The DFS must begin at every cell, not only the first. Look at the double loop at the bottom of the program. A loop that starts at r = 1 would skip the top row, and oath starts at the top row.

Third, check the path by hand. Find o at row 0, column 0. Its neighbour to the right is a at column 1. Its neighbour below that a is t. Below t is h. If any of those cells is # when the search reaches it, the cell was not restored after an earlier path.

Fourth, check the restoring line. The line board[r][c] = ch must run after the loop over the four neighbours, not inside it. If it sits inside the loop, the cell is restored too early, and a different path can reuse it in the middle of a word.

Fifth, check the early return. If node is None returns before the cell is marked, the cell stays unmarked, which is correct. If the mark happens before the check for None, a cell that is not on any word's path stays # for the rest of the search.

Work through the steps in this order. Most bugs show up in the first two, and the last three are the ones beginners miss.

What a missing restore does to the board

The restore line is easy to forget, and the program still runs without an error when it is missing. This section shows what goes wrong, using the same board and the same four words.

Remove the line board[r][c] = ch from search and run the program on the sample. The output is:

oath

The word eat has disappeared, and nothing crashed. The reason is the order of the search. The oath search starts at row 0, column 0, then walks to a at row 0, column 1, then t at row 1, column 1, then h at row 2, column 1. It marks each of those cells. Without the restore, they stay marked with # after the search returns.

When the loop later reaches row 1, column 3, the search for eat goes left to row 1, column 2, then tries to go left again to row 1, column 1, which now holds #. The check if ch == "#": return stops the search there, so eat is never found.

This is the trap of shared state. Several searches share one board. When one search changes a cell and does not change it back, every later search sees the change. The cure is always the same: whatever a call marks, the same call unmarks before it returns. The restore line is the second half of the mark, so the two lines must always come as a pair.

Variations and edge cases

Limit zero. A limit of zero means no words are wanted. The first check at the top of collect returns at once, so the result is empty, and that is correct.

A prefix that is a word. Suppose the store holds car and you ask for suggestions for car. The walk reaches the r node, the DFS adds car first, and then its children. So the word itself comes first in the list. Check your problem statement to see whether the word itself should be included.

A word longer than the board. A word of 10 letters cannot be traced on a 3 by 3 board, which has only 9 cells. The search simply never reaches a depth of 10 and never reports it. No special check is needed.

One cell with several words. The trie handles the case where the same letter starts several words. Each word has its own branch, and the search follows all of them at once.

The same word twice in the list. The word field is set once. The second copy overwrites the first with the same string, so it is still reported once.

An empty grid or no words. With no words, the trie root has no children, so every search returns at once. With no grid, the double loop does nothing. Both cases print nothing.

A board where every cell matches. Suppose the board is all a letters and the word list holds aaaa. The search finds many paths for the same word. The word field handles the repeat: the first path that reaches the end reports the word and clears the field, so later paths find nothing new. The output lists aaaa once. The search still explores every path, though, and that is the cost to watch on a large board of repeated letters.

Interview-style questions and real uses

Autocomplete is asked in interviews as "return the top k suggestions for a prefix." The answer has three steps: walk to the prefix node, run a DFS, and stop at k. Mention the order of the children, because it decides which words come first. If the ranking is by frequency rather than alphabetically, you need to store a score on each node and sort, which is a different design.

Word Search II is asked as "find all words in a board." The key is to say why one trie beats one search per word. The trie shares the work of common prefixes, and pruning cuts a path the moment no word can use it.

Real systems use the same ideas. A search box shows suggestions as you type, often from a trie or a similar prefix index. Spell checkers and word games use trie searches over a board or a rack of letters. Autocomplete on a phone keyboard walks a structure like this one for every key press, which is why it feels instant.

Practice

1. Which words does suggest("ca") return for the list car, card, care, cat, dog? Give the order too. Hint: four words, and the order is alphabetical.

Show the answer: It returns ['car', 'card', 'care', 'cat']. Those four words all start with ca, and the DFS visits them in alphabetical order.

2. Why is one trie better than searching the grid once for each word? Give one concrete example from the board in this lesson.

Show the answer: oath and oat both start with the same grid path from the top-left o to the a next to it. One trie search explores that path once, and then follows the branch for both words. A search per word would explore it twice.

3. Change suggest so it returns at most two words, without using a limit parameter. What does the code do before visiting each node?

Show the answer: Keep a counter of words added, or check len(out) at the top of collect. If the list already holds two words, return. The check at the top is what stops the walk early.

4. For the board in this lesson, does the program find pea? Explain by looking at the letters.

Show the answer: No. The board has no p, so no cell matches the first letter of pea, and the root's child for p is never reached on the board.

5. Explain what would go wrong if board[r][c] = ch were removed from the end of search.

Show the answer: Every cell visited once would stay #. Later paths could not use those cells, so words that need them are missed, and the output is incomplete. The program gives no error, which makes this bug hard to see. On the sample board the output is oath alone.

6. Count the nodes visited by collect for the prefix ca with no limit, in the example with car, card, care, cat and dog. Hint: the ca node itself counts.

Show the answer: Five: ca, car, card, care and cat. The dog branch is not under ca, so it is not counted.

7. Write the one-line check that makes collect stop at a limit of k, and say where it goes.

Show the answer: if len(out) >= limit: return, placed at the very top of collect before the node is visited.

8. Why does the DFS for autocomplete never walk into an empty branch?

Show the answer: Every trie node exists only because a stored word passes through it, so every node leads to at least one word. The DFS only reaches nodes that have a word below them.

9. Suppose the Word Search code set node.word on every node along the path, not only on the last one. What would the program print for the board in this lesson, and why?

Show the answer: It would print every prefix that is traceable and happens to be a prefix of a stored word, such as oa, oat and e, along with the real words. The word field must be set only on the last node of a word, because that node alone marks a complete word.

10. With a limit of k, give an upper bound on the nodes the autocomplete DFS visits below the prefix node. Then explain why it holds.

Show the answer: About k times M, where M is the length of the longest word. Every visited node lies on the path to a word that gets added, and each such path has at most M nodes below the prefix. There are at most k added words.

What you now know

• Autocomplete walks to the prefix node, then runs a DFS below it, collecting every word end.
• Children visited in slot order (a to z) give alphabetical output for free.
• A limit stops the DFS early, and the work is bounded by the limit times the longest word.
• Word Search II puts all words in one trie, runs one DFS per cell, and prunes when the trie has no matching child.
• Backtracking marks a cell on the path and restores it afterwards, so other paths can use it.
• Watch for the missing restore, the bounds check that Python does not enforce for negative indexes, and words reported twice.

In the next lesson the same trie is used on bits instead of letters, which turns a question about pairs of numbers into a walk of 31 steps.