Module assignment: word store

What you will build

You will write a small word store. It keeps a set of words, and it answers five kinds of command typed one per line: add a word, check whether a whole word is stored, count the stored words that start with some letters, suggest up to three words for some letters, and remove a word. Each command is answered with one line of output. The store is the same kind of trie you built in the first four lessons, so this assignment checks that you can join the pieces together: insert, search, starts_with, counting, the alphabetical walk, and remove.

The assignment is graded at three levels, described in the rubric below. A solution that does only the first two commands can earn a passing grade. A complete solution handles every command, every invalid input and every edge case in the list.

The commands

Each line of input holds one command and one argument, separated by a single space. The program reads lines until the input ends. In Python, a loop over sys.stdin does this. Each pass through the loop gets one line, including its newline character at the end, and line.split() breaks it into the command and the argument.

Command	Argument	Output
`add`	a word	`added`, or `already stored` if the word is already there
`has`	a word	`yes` if the word is stored as a whole word, otherwise `no`
`prefix`	letters	the number of different stored words that start with those letters
`suggest`	letters	up to three stored words that start with those letters, in alphabetical order, separated by single spaces, or `none`
`remove`	a word	`removed` if the word was stored, otherwise `not stored`

Every other command, such as find or delete, prints invalid. So does any argument that is not made of 1 to 20 lowercase letters from a to z. That covers capital letters, digits, spaces, punctuation and empty input. A line with no argument, or with more than one space-separated part, also prints invalid. Check the argument before you touch the store, so that an invalid argument never changes anything.

Two details are easy to miss. A word is stored as a whole word only when you added it. Asking has about a prefix of a stored word gives no, because the prefix is not a stored word. Also, prefix and suggest count and list words that start with the argument, and a stored word equal to the argument counts as starting with it.

Alphabetical order is the dictionary order used in the autocomplete lesson. The word car comes before card, and card comes before care, and care comes before cat, because r comes before t.

A sample to test against

Save the following lines in a file named sample1.txt, one command per line:

add car
add cat
add care
add car
has car
has ca
prefix ca
suggest c
suggest d
add Dog
remove car
has car
prefix c
remove dog
suggest c

Your program must print exactly these fifteen lines. To run a program on the file, use python3 main.py < sample1.txt. The < sends the file's contents to the program as if you had typed them.

added
added
added
already stored
yes
no
3
car care cat
none
invalid
removed
no
2
not stored
care cat

Work through a few of them before you run the program. The line prefix ca prints 3, not 2, because car, cat and care all start with ca at that point. The line suggest c prints car care cat, because care sorts before cat. The line add Dog prints invalid, because D is a capital letter. The line remove car prints removed, and after it the line prefix c prints 2, because only cat and care are left.

Here is a second, shorter test. Save it as sample2.txt and run it the same way. Its expected output is listed under it.

add zebra
has zebra
remove zebra
has zebra
remove zebra
add ab
add abc
prefix a
suggest a
add 9x
foo bar

added
yes
removed
no
not stored
added
added
2
ab abc
invalid
invalid

The last line, foo bar, has two parts, so it prints invalid. The line add 9x is an invalid argument, because it contains a digit.

The requirements checklist

Tick each item as you finish it. Every item is a behaviour you can test with the sample above or with your own inputs.

Grading rubric

The grade depends on how many of the commands work correctly on hidden tests. Each level below is a full description of what a solution at that level does.

60% (passing). The program handles add, has and invalid input. It prints added, already stored, yes, no and invalid correctly. The prefix, suggest and remove commands may print invalid or nothing useful. This level shows that you can insert and search and that you keep invalid input out of the store.

80% (good). The program also handles prefix and suggest. Counts are correct, the three-word limit is respected, and the alphabetical order is right, including the case where a word is a prefix of another word. remove may still be missing. This level shows that you can walk to a prefix node, count the words under it, and collect the words in order.

100% (complete). Every command works, including remove. After a removal, the counts and suggestions no longer include the removed word, and a later add of the same word works again. The code is organised as a class, the invalid input check runs first, and the answers come from the trie structure, not from a list of all words that is scanned on each command.

To earn the top grade, test the removal rules yourself. Add two words that share a long prefix, remove one of them, and check that prefix and suggest still show the other. Then remove the second one and check that the prefix is gone. Those two checks catch most mistakes in the counting logic.

Starter code

Start from this program. It reads each command and its argument, and prints a placeholder line for each one. Replace the placeholder with your own work, one command at a time. Build the store first, then add each command.

import sys
 
for line in sys.stdin:
    command, arg = line.split()
    # TODO: decide what each command does and print one line for it.
    print("todo:", command, arg)

Run it with python3 main.py < sample1.txt. You should see fifteen todo lines, one for each command, with the command and its argument echoed back. That confirms that the reading loop works before you write any logic.

Notice what the starter does not do. It does not check the argument, and it does not store anything. Both come next. Writing the reading loop first means a broken reader is never blamed for a broken trie. The starter also crashes on a blank line, because line.split() gives an empty list and the unpacking fails. The milestone programs below skip blank lines, so you can copy that check when you need it.

Milestone one: a 60 percent solution

This is a complete program for the first milestone. It handles add, has and the invalid check, and prints invalid for every other command. Read it closely, because the second and third milestones build on it.

import sys
 
 
class Node:
    def __init__(self):
        self.children = [None] * 26
        self.is_word = False
 
 
root = Node()
 
 
def is_valid(s):
    """A word is 1 to 20 lowercase letters from a to z."""
    return 1 <= len(s) <= 20 and all("a" <= ch <= "z" for ch in s)
 
 
def add(word):
    node = root
    for ch in word:
        i = ord(ch) - ord("a")
        if node.children[i] is None:
            node.children[i] = Node()
        node = node.children[i]
    if node.is_word:
        return False
    node.is_word = True
    return True
 
 
def has(word):
    node = root
    for ch in word:
        node = node.children[ord(ch) - ord("a")]
        if node is None:
            return False
    return node.is_word
 
 
for line in sys.stdin:
    if not line.strip():
        continue
    parts = line.split()
    if len(parts) != 2:
        print("invalid")
        continue
    command, arg = parts
    if not is_valid(arg):
        print("invalid")
        continue
    match command:
        case "add":
            print("added" if add(arg) else "already stored")
        case "has":
            print("yes" if has(arg) else "no")
        case _:
            print("invalid")

Run this with sample1.txt. The first six lines of the output match the expected output: added, added, added, already stored, yes, no. The lines for prefix, suggest and remove print invalid in this version, which is expected at this level. The line add Dog already prints invalid, which is correct.

Look at add closely. It walks to the end of the word, creating nodes on the way, then checks the flag. The check comes after the walk, because the node exists only when the walk is finished. If the flag is already set, the word was stored before, so the function returns False and leaves the trie as it was. Creating the missing nodes before that check is harmless, because a word that is already stored has all of its nodes already.

Milestone two: an 80 percent solution

The second milestone adds prefix and suggest. It keeps the same add and has, and it adds a counter on every node. The counter is the number of stored words that pass through the node, as you saw in the second lesson. A walk to a prefix node then returns the count in one step. The suggest command reuses the alphabetical walk from the third lesson, stopped after three words.

import sys
 
 
class Node:
    def __init__(self):
        self.children = [None] * 26
        self.is_word = False
        self.pass_count = 0
 
 
root = Node()
 
 
def is_valid(s):
    return 1 <= len(s) <= 20 and all("a" <= ch <= "z" for ch in s)
 
 
def walk(text):
    node = root
    for ch in text:
        node = node.children[ord(ch) - ord("a")]
        if node is None:
            return None
    return node
 
 
def add(word):
    existing = walk(word)
    if existing is not None and existing.is_word:
        return False
    node = root
    node.pass_count += 1
    for ch in word:
        i = ord(ch) - ord("a")
        if node.children[i] is None:
            node.children[i] = Node()
        node = node.children[i]
        node.pass_count += 1
    node.is_word = True
    return True
 
 
def has(word):
    node = walk(word)
    return node is not None and node.is_word
 
 
def count_prefix(prefix):
    node = walk(prefix)
    return 0 if node is None else node.pass_count
 
 
def collect(node, path, out, limit):
    if len(out) >= limit:
        return
    if node.is_word:
        out.append("".join(path))
    for i, child in enumerate(node.children):
        if child is not None:
            path.append(chr(ord("a") + i))
            collect(child, path, out, limit)
            path.pop()
 
 
def suggest(prefix, limit):
    out = []
    node = walk(prefix)
    if node is not None:
        collect(node, list(prefix), out, limit)
    return out
 
 
for line in sys.stdin:
    if not line.strip():
        continue
    parts = line.split()
    if len(parts) != 2:
        print("invalid")
        continue
    command, arg = parts
    if not is_valid(arg):
        print("invalid")
        continue
    match command:
        case "add":
            print("added" if add(arg) else "already stored")
        case "has":
            print("yes" if has(arg) else "no")
        case "prefix":
            print(count_prefix(arg))
        case "suggest":
            words = suggest(arg, 3)
            print(" ".join(words) if words else "none")
        case _:
            print("invalid")

This version gets the sample right up to the remove lines. It prints invalid for remove car and remove dog, which is expected at this level. Because car was never removed, the later lines are also different from the expected output: has car prints yes, prefix c prints 3 instead of 2, and suggest c prints car care cat. Those three differences all come from the missing remove. The line prefix ca must print 3, and the line suggest c must print car care cat before any removal.

Milestone three: the full solution

The last milestone adds remove. The rule is the one from the second lesson. Clear the flag at the word's last node, subtract one from the counter of every node on the path, and cut a link when a counter reaches zero. Because the counters are correct after a removal, both prefix and suggest update on their own, with no extra work.

The reference solution is below. Try to finish the program yourself first. Use the hints if you get stuck, and open the solution only after you have tried. The checklist above tells you what the program must do.

Hint 1. Keep the store in one class, with add, has, count_prefix, suggest and remove as methods. Let the main part of the program call them and print the answer. This makes it easy to test each method alone.

Hint 2. For remove, first check has(word). If the word is not stored, return False and change nothing. Only then start the walk that decreases counters.

Hint 3. The root counter counts every stored word, so decrease it once at the start of remove. Then walk the word, and at each step decrease the child's counter. If a child's counter becomes zero, set the link from its parent to None, and return. The nodes below that child are no longer reachable, so you do not need to update them.

Hint 4. After the walk finishes without cutting a link, the last node is still in the trie, so clear its is_word flag. The word is gone, but a longer word may still pass through that node.

Hint 5. Put the invalid check at the top of the command loop, before any match. Then no command, valid or not, can change the store when its argument is invalid.

Hint 6. Test with sample2.txt. Its remove zebra line is the case where the word is stored, and the next remove zebra is the case where it is not. The line has zebra after the first removal must print no.

The reference solution

The solution below is complete and handles every command. Use it to check your work after you have tried the problem yourself. It is also a model of the layout the rubric asks for: one class holding the store, and a main function that only reads and prints. The last two lines make main run when you start the file, and not when another file imports it.

The program below is the reference. Read it after your own attempt.

import sys
 
 
class Node:
    def __init__(self):
        self.children = [None] * 26
        self.is_word = False
        self.pass_count = 0  # number of stored words that pass through this node
 
 
class WordStore:
    def __init__(self):
        self.root = Node()
 
    def _walk(self, text):
        node = self.root
        for ch in text:
            node = node.children[ord(ch) - ord("a")]
            if node is None:
                return None
        return node
 
    def has(self, word):
        node = self._walk(word)
        return node is not None and node.is_word
 
    def count_prefix(self, prefix):
        node = self._walk(prefix)
        return 0 if node is None else node.pass_count
 
    def add(self, word):
        if self.has(word):
            return False
        node = self.root
        node.pass_count += 1
        for ch in word:
            i = ord(ch) - ord("a")
            if node.children[i] is None:
                node.children[i] = Node()
            node = node.children[i]
            node.pass_count += 1
        node.is_word = True
        return True
 
    def remove(self, word):
        if not self.has(word):
            return False
        node = self.root
        node.pass_count -= 1
        for ch in word:
            i = ord(ch) - ord("a")
            child = node.children[i]
            child.pass_count -= 1
            if child.pass_count == 0:
                node.children[i] = None
                return True
            node = child
        node.is_word = False
        return True
 
    def _collect(self, node, path, out, limit):
        if len(out) >= limit:
            return
        if node.is_word:
            out.append("".join(path))
        for i, child in enumerate(node.children):
            if child is not None:
                path.append(chr(ord("a") + i))
                self._collect(child, path, out, limit)
                path.pop()
 
    def suggest(self, prefix, limit):
        out = []
        node = self._walk(prefix)
        if node is not None:
            self._collect(node, list(prefix), out, limit)
        return out
 
 
def is_valid(s):
    return 1 <= len(s) <= 20 and all("a" <= ch <= "z" for ch in s)
 
 
def main():
    store = WordStore()
    for line in sys.stdin:
        if not line.strip():
            continue
        parts = line.split()
        if len(parts) != 2:
            print("invalid")
            continue
        command, arg = parts
        if not is_valid(arg):
            print("invalid")
            continue
        match command:
            case "add":
                print("added" if store.add(arg) else "already stored")
            case "has":
                print("yes" if store.has(arg) else "no")
            case "prefix":
                print(store.count_prefix(arg))
            case "suggest":
                words = store.suggest(arg, 3)
                print(" ".join(words) if words else "none")
            case "remove":
                print("removed" if store.remove(arg) else "not stored")
            case _:
                print("invalid")
 
 
if __name__ == "__main__":
    main()

Run it with sample1.txt and sample2.txt. The output must match the expected lines exactly.

Walking through the reference

The reference keeps all of the store inside WordStore. The main function does three things for each command: it reads the command and argument, it runs the invalid check, and it picks the right branch of the match. This split means the store can be tested without typing commands, and each method has one job.

The remove method is the only new piece. It calls has first, which handles the case of a word that is not stored. Then it decreases the root counter once, and walks down the word, decreasing each counter. When a counter reaches zero, the link from its parent is cut and the method returns True. This is correct because a zero counter means no stored word passes through that node. The nodes below it are no longer reachable from the root, so they do not need to be updated.

If the walk finishes without cutting any link, the last node still has a counter above zero. Its word flag is cleared, because the word is no longer stored. Longer words that pass through the same node keep their counters, so they still count in prefix and still appear in suggest.

Why the cost stays low

Each command walks its argument once, which takes time proportional to the argument's length. suggest also visits up to three words, each of at most 20 letters, so its extra work is bounded by a small constant. None of the commands scans the list of all stored words. So a store with a hundred thousand words answers each command in about the same time as one with ten words, when the argument is the same length.

Compare this with a list that you scan on every command. Each prefix would check every word, so the time would grow with the store. The trie, with its counters, avoids that. The note at the end can mention this difference.

Common mistakes in this assignment

Forgetting the invalid check for remove. The check belongs before the match, so that remove Dog prints invalid and not not stored.

Using isalpha() for the letter check. It looks like a shortcut, but Python's str.isalpha() returns True for letters from other alphabets. For example, "café".isalpha() is True, and the word would then be stored as if it were valid. The check "a" <= ch <= "z" in the reference accepts only the 26 letters. Run print("café".isalpha()) to see the difference for yourself.

Changing the store on an invalid command. If add runs before the check, a 21-letter word such as abcdefghijklmnopqrstu is stored. Check first.

Counting prefix from the wrong node. prefix must read the counter of the prefix node. Reading the root counter gives the total for every prefix, which is wrong.

Printing none for a valid empty result from prefix. prefix prints a number, even zero. Only suggest prints none.

Suggesting a removed word. If remove clears the flag but leaves the counters alone, the counters are wrong and prefix overcounts. If it updates the counters but never cuts a link at zero, suggest can still reach a node with no word below it. Test both.

Leaving suggest unsorted. The walk gives alphabetical order only when it visits children from a to z. Do not reorder the children, and do not sort the output on a whim, because that hides an error in the walk.

Forgetting to handle words that are prefixes of other words. car and card are both stored. Removing car must leave card in place, and has car must then print no, while prefix car still counts card.

Testing your solution

Test in three rounds, from small to large. First run the sample and compare line by line with the expected output. If one line differs, find the first line that is wrong, because each later line depends on the state the earlier commands left behind. A command such as diff on the output file and a text file of the expected lines does this comparison in one step.

Second, write your own short input that tests one rule at a time. Add two words with a shared prefix, check prefix on the shared prefix, remove one word, and check again. Add the removed word back and check that added comes out and prefix returns to its old value. Try an invalid argument between valid commands and check that the store does not change.

Third, try a larger input made by hand. Type twenty or so commands with overlapping prefixes into a file. Predict each answer on paper before you run the program. A wrong prediction is useful: it tells you whether the bug is in your prediction or in the code.

What to hand in

Submit the complete main.py file, with the WordStore class, the invalid check and the main function. Include a comment above each method that says what it does in one sentence. Keep the code as you wrote it, without the reference solution, so the grader sees your own design.

Along with the file, write a short note of at most ten lines. Say which level of the rubric you think you reached, which command you found hardest, and which test caught your most serious bug. The note is part of the grade, because it shows that you can test and explain your work, not only write it.

What this assignment checks

The assignment uses every part of the module. The first lesson's node design is the base. The second lesson's insert, search, starts_with, counters and remove are the core of the store. The third lesson's walk and limit give the suggestions. The fourth lesson's greedy walk is not needed here, but the note asks which test caught your most serious bug.

If you finish with a score of 100 percent, you have a working trie with counts, a correct removal and a clean interface. That is the foundation for the next module, where a trie is one of several structures used to answer questions about text.