Tries

Problem¶

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.

So:

Nodes represent prefix states.
Edges are labeled with characters.

In the implementation below, this line shows the edge-label view:

node.children[ch] = TrieNode()

ch is the edge label.
TrieNode() is the destination node.
is_word belongs to the node because the node represents the prefix reached so far.

Example¶

insert: "cat", "car"

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

Or more precisely:

root
 │  c
 └───── .
        │  a
        └───── .
               │  t
               └───── .*
               │  r
               └───── .*

where

. is the node.
c, a, t, r are edge labels.
* marks a node where a full word ends.

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