Tries

Question¶

Define tries and provide an implementation in Python.

Definition¶

Tree for strings where each edge is a character and each node stores whether a word ends there.

Example¶

insert: "cat", "car"

root
 └─ c
    └─ a
       ├─ t*
       └─ r*

Key trick¶

Shared prefixes give fast prefix queries.

Trap¶

Distinguish prefix exists from full word exists.

Use¶

Autocomplete, dictionary, prefix search.

Do not use¶

When memory is tight and simple hashing is enough.

Implementation¶

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

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

    def insert(self, word: str) -> None:
        node = self.root
        for ch in word:
            node = node.children.setdefault(ch, TrieNode())
        node.is_word = True

    def search(self, word: str) -> bool:
        node = self.root
        for ch in word:
            if ch not in node.children:
                return False
            node = node.children[ch]
        return node.is_word

    def startswith(self, prefix: str) -> bool:
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return False
            node = node.children[ch]
        return True

Pytest¶

import pytest

@pytest.mark.parametrize(
    "words, search_word, prefix, expected_search, expected_prefix",
    [
        (["cat", "car"], "cat", "ca", True, True),
        (["cat", "car"], "cap", "ca", False, True),
        (["dog"], "do", "do", False, True),
        ([], "a", "a", False, False),
    ],
)
def test_trie(words, search_word, prefix, expected_search, expected_prefix):
    trie = Trie()
    for w in words:
        trie.insert(w)

    assert trie.search(search_word) == expected_search
    assert trie.startswith(prefix) == expected_prefix