{"id":"dsa-trees-tries","slug":"tries","title":"Challenge: tries","text":"A search box suggests words as you type: after \"ca\", it offers \"cab\", \"cake\", \"car\" and \"card\". With the words in a list, every question scans all of them with `startswith`: O(n·m) for n words and a prefix of m letters, again for each letter typed. A **trie** (from re*trie*val, said \"try\") stores the words letter by letter in a tree, so that words with the same beginning share it. Each node stands for a prefix. Its children are a dict from a letter to the next node, and a flag marks the nodes where a whole word ends. Here are \"cab\", \"cake\", \"car\", \"card\" and \"dog\", with a `*` where a word ends:\n\n```text\n(root)\n +-- c\n |   +-- a\n |       +-- b *\n |       +-- k\n |       |   +-- e *\n |       +-- r *\n |           +-- d *\n +-- d\n     +-- o\n         +-- g *\n```\n\n```python\nclass TrieNode:\n    def __init__(self):\n        self.children = {}\n        self.is_word = False\n\n\nclass Trie:\n    def __init__(self):\n        self.root = TrieNode()\n\n    def insert(self, word):\n        node = self.root\n        for letter in word:\n            if letter not in node.children:\n                node.children[letter] = TrieNode()\n            node = node.children[letter]\n        node.is_word = True\n\n    def find(self, prefix):\n        node = self.root\n        for letter in prefix:\n            if letter not in node.children:\n                return None\n            node = node.children[letter]\n        return node\n\n    def contains(self, word):\n        node = self.find(word)\n        return node is not None and node.is_word\n\n    def complete(self, prefix):\n        words = []\n        node = self.find(prefix)\n        if node is not None:\n            self.collect(node, prefix, words)\n        return words\n\n    def collect(self, node, text, words):\n        if node.is_word:\n            words.append(text)\n        for letter in sorted(node.children):\n            self.collect(node.children[letter], text + letter, words)\n\n\ntrie = Trie()\nfor word in [\"car\", \"dog\", \"card\", \"cake\", \"cab\"]:\n    trie.insert(word)\nprint(trie.contains(\"car\"), trie.contains(\"ca\"), trie.contains(\"cart\"))\nprint(trie.complete(\"ca\"))\nprint(trie.complete(\"do\"), trie.complete(\"x\"))\n```\n\n```output\nTrue False False\n['cab', 'cake', 'car', 'card']\n['dog'] []\n```\n\n`insert` walks down the word, making the nodes that are missing, and marks the last one. `find` walks down a prefix one letter at a time and returns its node, or `None` when the path breaks off. \"ca\" is in the trie as a prefix, but no word ends there, which is why `contains` checks the flag.\n\n### Autocomplete in sorted order\n\nEvery word that starts with a prefix lies below the prefix's node. So `complete` finds that node, and `collect` visits everything under it in pre-order: a node's own word first, then its children. Two things make the words come out sorted, with no sorting of the words at the end. A word comes before the longer words that start with it (\"car\" before \"card\"), and the children are visited in alphabetical order.\n\nThat `sorted` matters. A dict keeps its keys in the order they were added, so a `collect` that loops over `node.children` as it is gives the words in an order that depends on the inserts:\n\n```python\nclass TrieNode:\n    def __init__(self):\n        self.children = {}\n        self.is_word = False\n\n\ndef insert(root, word):\n    node = root\n    for letter in word:\n        if letter not in node.children:\n            node.children[letter] = TrieNode()\n        node = node.children[letter]\n    node.is_word = True\n\n\ndef collect(node, text, words):\n    if node.is_word:\n        words.append(text)\n    for letter in node.children:\n        collect(node.children[letter], text + letter, words)\n\n\nroot = TrieNode()\nfor word in [\"car\", \"dog\", \"card\", \"cake\", \"cab\"]:\n    insert(root, word)\nwords = []\ncollect(root, \"\", words)\nprint(words)\n```\n\n```output\n['car', 'card', 'cake', 'cab', 'dog']\n```\n\nBelow \"ca\", the letters were added in the order r, k, b, because \"car\" came before \"cake\" and \"cab\", so \"car\" and \"card\" come out first. The same five words inserted in another order would come out in another order again. A test with words inserted in alphabetical order hides this mistake, which is why good tests mix them up.\n\nThe recursion in `collect` goes one level deeper per letter, so its depth is the length of the longest word: no danger of a `RecursionError`.\n\n### What it costs\n\n`insert`, `find` and `contains` take one dict step per letter: O(m) for a word or a prefix of m letters, whether the trie holds ten words or a million. `complete` costs O(m) to find the node, then a visit to each node below it. Each visit builds its node's text with `text + letter`, which copies every letter so far. So with k nodes below the prefix and words of at most L letters, `collect` costs O(k·L): close to O(k) for short words, but in a trie of one word of 500 letters, `complete(\"\")` copies over 125,000 letters. A version that keeps one list of letters, appends a letter on the way down, pops it on the way back up and calls `\"\".join` only where a word ends, costs O(k) plus the letters of the words it returns. The memory is one node per different prefix: at most the total number of letters in all the words, and less where words share their beginnings.\n\nA trie can count, too. If every node keeps a `count` of the words that pass through it (`insert` adds 1 to the root and to each node on its way down), the number of words that start with a prefix is the `count` of the prefix's node: one walk of m letters instead of a scan of every word. The root is the node of the empty prefix `\"\"`, which every word starts with, so its `count` is the number of words.\n","task":"A search engine shortens each word to the shortest **stem** it knows that the word starts with, so that \"raining\" and \"rainbows\" both become \"rain\". The stems are in a trie. Fill in the blanks of `shortest_stem`: it walks down the word one letter at a time, and at the first node where a stem ends, it returns the part of the word read so far, `word[0]` to `word[i]`. A word that starts with no stem stays as it is. The program should print:\n\n```\n['play', 'cat', 'saw', 'rain', 'in', 'the', 'rain']\nca rain\n```\n","kind":"fill","seed":"class TrieNode:\n    def __init__(self):\n        self.children = {}\n        self.is_word = False\n\n\nclass Trie:\n    def __init__(self):\n        self.root = TrieNode()\n\n    def insert(self, word):\n        node = self.root\n        for letter in word:\n            if letter not in node.children:\n                node.children[letter] = TrieNode()\n            node = node.children[letter]\n        node.is_word = True\n\n    def shortest_stem(self, word):\n        node = self.root\n        for i in range(len(word)):\n            letter = word[i]\n            if letter not in node.children:\n                return word\n            node = node.children[letter]\n            if [[node.is_word‖node.is_word == True‖node.is_word is True]]:\n                return word[:[[i + 1‖1 + i]]]\n        return word\n\n\ntrie = Trie()\nfor stem in [\"rainbow\", \"play\", \"rain\", \"cat\"]:\n    trie.insert(stem)\nwords = \"playful cats saw rainbows in the rain\".split()\nprint([trie.shortest_stem(word) for word in words])\nprint(trie.shortest_stem(\"ca\"), trie.shortest_stem(\"raindrop\"))\n","solution":"class TrieNode:\n    def __init__(self):\n        self.children = {}\n        self.is_word = False\n\n\nclass Trie:\n    def __init__(self):\n        self.root = TrieNode()\n\n    def insert(self, word):\n        node = self.root\n        for letter in word:\n            if letter not in node.children:\n                node.children[letter] = TrieNode()\n            node = node.children[letter]\n        node.is_word = True\n\n    def shortest_stem(self, word):\n        node = self.root\n        for i in range(len(word)):\n            letter = word[i]\n            if letter not in node.children:\n                return word\n            node = node.children[letter]\n            if node.is_word:\n                return word[:i + 1]\n        return word\n\n\ntrie = Trie()\nfor stem in [\"rainbow\", \"play\", \"rain\", \"cat\"]:\n    trie.insert(stem)\nwords = \"playful cats saw rainbows in the rain\".split()\nprint([trie.shortest_stem(word) for word in words])\nprint(trie.shortest_stem(\"ca\"), trie.shortest_stem(\"raindrop\"))\n","hints":["A stem ends at a node whose `is_word` flag is set. Check it right after stepping into the node for `letter`: the first such node is the shortest stem, so \"rainbows\" stops at \"rain\", not at \"rainbow\".","The first blank is `node.is_word`. The letters read so far are those at the indexes 0 to `i`, and a slice stops just before its end index: `word[:i + 1]`."],"tests":[{"name":"Output","stdin":"","expect":"['play', 'cat', 'saw', 'rain', 'in', 'the', 'rain']\nca rain","hidden":false}],"require":[],"forbid":[],"more":[{"task":"Write `count_with_prefix(words, prefixes)` that returns a list with, for each prefix, the number of words that start with it. A word that is in the list twice counts twice, and every word starts with the empty prefix `\"\"`. The hidden test has 60,000 words and 60,000 prefixes, so a scan of all the words for each prefix is too slow: build a trie whose nodes count the words that pass through them.\n\nFor example, this program:\n\n```python\nwords = [\"car\", \"card\", \"care\", \"cat\", \"dog\", \"car\"]\nprint(count_with_prefix(words, [\"car\", \"ca\", \"d\", \"cart\", \"x\"]))\n```\n\nprints:\n\n```\n[4, 5, 1, 0, 0]\n```\n","kind":"code","seed":"def count_with_prefix(words, prefixes):\n    pass\n\n\nif __name__ == \"__main__\":\n    words = [\"car\", \"card\", \"care\", \"cat\", \"dog\", \"car\"]\n    print(count_with_prefix(words, [\"car\", \"ca\", \"d\", \"cart\", \"x\"]))\n","solution":"class TrieNode:\n    def __init__(self):\n        self.children = {}\n        self.count = 0\n\n\ndef count_with_prefix(words, prefixes):\n    root = TrieNode()\n    for word in words:\n        node = root\n        node.count += 1\n        for letter in word:\n            if letter not in node.children:\n                node.children[letter] = TrieNode()\n            node = node.children[letter]\n            node.count += 1\n    result = []\n    for prefix in prefixes:\n        node = root\n        for letter in prefix:\n            if letter not in node.children:\n                node = None\n                break\n            node = node.children[letter]\n        if node is None:\n            result.append(0)\n        else:\n            result.append(node.count)\n    return result\n\n\nif __name__ == \"__main__\":\n    words = [\"car\", \"card\", \"care\", \"cat\", \"dog\", \"car\"]\n    print(count_with_prefix(words, [\"car\", \"ca\", \"d\", \"cart\", \"x\"]))\n","hints":["Make a node class with a `children` dict and a `count`. Insert every word as the lesson's `insert` does, and add 1 to the `count` of the root and of each node you step into.","Then the count of a prefix's node is the number of words that pass through it, which are the words that start with the prefix. The root's count is the answer for the empty prefix `\"\"`: every word starts with it.","For each prefix, walk down from the root letter by letter. If a letter is missing, no word starts with the prefix: its answer is 0. Otherwise, the answer is the count of the node where the walk ends."],"tests":[{"name":"Test 1","stdin":"","call":"words = [\"car\", \"card\", \"care\", \"cat\", \"dog\", \"car\"]\nprint(count_with_prefix(words, [\"car\", \"ca\", \"d\", \"cart\", \"x\"]))\n","expect":"[4, 5, 1, 0, 0]","hidden":false},{"name":"count_with_prefix([], [\"a\", \"b\"])","stdin":"","call":"count_with_prefix([], [\"a\", \"b\"])\n","expect":"[0, 0]","hidden":true,"returns":true},{"name":"count_with_prefix([\"apple\", \"app\", \"apply\", \"ape\"], [\"app\", \"apple\", \"apples\", \"a\", \"p\"])","stdin":"","call":"count_with_prefix([\"apple\", \"app\", \"apply\", \"ape\"], [\"app\", \"apple\", \"apples\", \"a\", \"p\"])\n","expect":"[3, 1, 0, 4, 0]","hidden":true,"returns":true},{"name":"count_with_prefix([\"sun\", \"sun\", \"sunny\"], [\"sun\", \"sunn\", \"s\"])","stdin":"","call":"count_with_prefix([\"sun\", \"sun\", \"sunny\"], [\"sun\", \"sunn\", \"s\"])\n","expect":"[3, 1, 3]","hidden":true,"returns":true},{"name":"count_with_prefix([\"a\", \"ab\", \"abc\"], [\"abc\", \"abcd\", \"b\"])","stdin":"","call":"count_with_prefix([\"a\", \"ab\", \"abc\"], [\"abc\", \"abcd\", \"b\"])\n","expect":"[1, 0, 0]","hidden":true,"returns":true},{"name":"count_with_prefix([\"car\", \"card\", \"dog\"], [\"\", \"ca\"])","stdin":"","call":"count_with_prefix([\"car\", \"card\", \"dog\"], [\"\", \"ca\"])\n","expect":"[3, 2]","hidden":true,"returns":true},{"name":"Test 7","stdin":"","call":"words = [str(i * 7919 % 10_007) for i in range(60_000)]\ncounts = count_with_prefix(words, [str(i * 37 % 100_003) for i in range(60_000)])\nprint(len(counts), max(counts), sum(counts))\n","expect":"60000 6705 147777","hidden":true,"big":true}],"require":[],"forbid":[]},{"task":"A word game lets players search with a dot for any one letter: \"c.t\" matches \"cat\", \"cot\" and \"cut\", but not \"cart\" (a dot is exactly one letter). Add a method `matches(self, pattern)` to the `Trie` class that returns the words of the trie that match `pattern`, in alphabetical order.\n\nWalk down the trie along the pattern. A letter follows its one child, if there is one. A dot follows every child, in alphabetical order. When the whole pattern is used up, the node you're at is a match if a word ends there. A helper method that does this recursively, given the node, the pattern and the text so far, keeps it short.\n\nFor example, this program:\n\n```python\ntrie = Trie()\nfor word in [\"cat\", \"cot\", \"cut\", \"cart\", \"dog\", \"dot\", \"cab\"]:\n    trie.insert(word)\nprint(trie.matches(\"c.t\"))\nprint(trie.matches(\"do.\"), trie.matches(\"x..\"))\n```\n\nprints:\n\n```\n['cat', 'cot', 'cut']\n['dog', 'dot'] []\n```\n","kind":"code","seed":"class TrieNode:\n    def __init__(self):\n        self.children = {}\n        self.is_word = False\n\n\nclass Trie:\n    def __init__(self):\n        self.root = TrieNode()\n\n    def insert(self, word):\n        node = self.root\n        for letter in word:\n            if letter not in node.children:\n                node.children[letter] = TrieNode()\n            node = node.children[letter]\n        node.is_word = True\n\n    def matches(self, pattern):\n        pass\n\n\nif __name__ == \"__main__\":\n    trie = Trie()\n    for word in [\"cat\", \"cot\", \"cut\", \"cart\", \"dog\", \"dot\", \"cab\"]:\n        trie.insert(word)\n    print(trie.matches(\"c.t\"))\n    print(trie.matches(\"do.\"), trie.matches(\"x..\"))\n","solution":"class TrieNode:\n    def __init__(self):\n        self.children = {}\n        self.is_word = False\n\n\nclass Trie:\n    def __init__(self):\n        self.root = TrieNode()\n\n    def insert(self, word):\n        node = self.root\n        for letter in word:\n            if letter not in node.children:\n                node.children[letter] = TrieNode()\n            node = node.children[letter]\n        node.is_word = True\n\n    def matches(self, pattern):\n        found = []\n        self.match(self.root, pattern, \"\", found)\n        return found\n\n    def match(self, node, pattern, text, found):\n        if len(text) == len(pattern):\n            if node.is_word:\n                found.append(text)\n            return\n        letter = pattern[len(text)]\n        if letter == \".\":\n            for child in sorted(node.children):\n                self.match(node.children[child], pattern, text + child, found)\n        elif letter in node.children:\n            self.match(node.children[letter], pattern, text + letter, found)\n\n\nif __name__ == \"__main__\":\n    trie = Trie()\n    for word in [\"cat\", \"cot\", \"cut\", \"cart\", \"dog\", \"dot\", \"cab\"]:\n        trie.insert(word)\n    print(trie.matches(\"c.t\"))\n    print(trie.matches(\"do.\"), trie.matches(\"x..\"))\n","hints":["Write a helper method `match(self, node, pattern, text, found)`: `text` is the part of the pattern matched so far, so `len(text)` tells which letter of the pattern comes next. `matches` calls it with the root, an empty text and an empty list.","The base case: when `len(text) == len(pattern)`, append `text` to `found` if `node.is_word`, and return.","Otherwise, look at `letter = pattern[len(text)]`. For a dot, call the helper for every child in `sorted(node.children)`, with that child's letter added to `text`. For any other letter, call it only for that letter's child, if there is one."],"tests":[{"name":"Test 1","stdin":"","call":"trie = Trie()\nfor word in [\"cat\", \"cot\", \"cut\", \"cart\", \"dog\", \"dot\", \"cab\"]:\n    trie.insert(word)\nprint(trie.matches(\"c.t\"))\nprint(trie.matches(\"do.\"), trie.matches(\"x..\"))\n","expect":"['cat', 'cot', 'cut']\n['dog', 'dot'] []","hidden":false},{"name":"Test 2","stdin":"","call":"trie = Trie()\nfor word in [\"cat\", \"cot\", \"cut\", \"cart\", \"dog\", \"dot\", \"cab\"]:\n    trie.insert(word)\nprint(trie.matches(\"...\"))\nprint(trie.matches(\"c..t\"), trie.matches(\".a.\"), trie.matches(\"ca\"))\n","expect":"['cab', 'cat', 'cot', 'cut', 'dog', 'dot']\n['cart'] ['cab', 'cat'] []","hidden":true},{"name":"Test 3","stdin":"","call":"trie = Trie()\nfor word in [\"moon\", \"mood\", \"moan\", \"mean\", \"men\"]:\n    trie.insert(word)\nprint(trie.matches(\"m..n\"), trie.matches(\"mo.d\"), trie.matches(\"m.n\"), trie.matches(\"....\"))\n","expect":"['mean', 'moan', 'moon'] ['mood'] ['men'] ['mean', 'moan', 'mood', 'moon']","hidden":true},{"name":"Test 4","stdin":"","call":"trie = Trie()\nprint(trie.matches(\"a.\"))\n","expect":"[]","hidden":true}],"require":[],"forbid":[]}]}