Build a Trie for AutocompleteStep 5 of 5

Counting

12345

Step 5: Counting

Finally, some analytics. Count words in the whole trie or under any prefix, handy for ranking autocomplete results ("1,204 matches").

  • count(): the total number of words stored
  • count_prefix(prefix): how many words start with prefix (0 if the prefix isn't present)

Both are a depth-first walk that tallies every node whose is_word is set. count_prefix just starts from the prefix node instead of the root.

Hints
Starter code
class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_word = False


class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.is_word = True

    def _find(self, prefix):
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return None
            node = node.children[ch]
        return node

    def _count_from(self, node):
        """Count word-ends in the subtree rooted at node."""
        # TODO: total = 1 if node.is_word else 0; add _count_from of each child.
        pass

    def count(self):
        """Total number of words in the trie."""
        # TODO: return self._count_from(self.root)
        pass

    def count_prefix(self, prefix):
        """Number of words that start with prefix (0 if absent)."""
        # TODO: node = self._find(prefix); return 0 if None else self._count_from(node)
        pass
Write it on a laptop

The editor and the test runner need a wide screen. Open this page on a desktop browser and your progress will be waiting.

PreviousAll challenges