Skip to content

Trie

Problem

Implement a prefix tree that supports insert, search, prefix matching, delete, and autocomplete operations on a set of strings.

Approach

A tree where each node represents a character. Paths from the root to nodes marked as word-endings form the stored words. Children are stored in a dict keyed by character for O(1) branching.

When to Use

Autocomplete, spell checking, IP routing, prefix matching. ASI relevance: airport code lookup (e.g., "BO" -> ["BOS", "BOI", "BOG"]), flight ID prefix search, geospatial name indexing.

Complexity
Time O(m) per insert/search/starts_with/delete where m = word length.
Space O(N * m) where N = number of words, m = average word length.

Implementation

trie

Trie (prefix tree) for efficient string operations.

Problem

Implement a prefix tree that supports insert, search, prefix matching, delete, and autocomplete operations on a set of strings.

Approach

A tree where each node represents a character. Paths from the root to nodes marked as word-endings form the stored words. Children are stored in a dict keyed by character for O(1) branching.

When to use

Autocomplete, spell checking, IP routing, prefix matching. ASI relevance: airport code lookup (e.g., "BO" -> ["BOS", "BOI", "BOG"]), flight ID prefix search, geospatial name indexing.

Complexity

Time: O(m) per insert/search/starts_with/delete where m = word length. O(p + k) for autocomplete where p = prefix length, k = total characters in matching subtree. Space: O(N * m) where N = number of words, m = average word length.

Trie

Prefix tree supporting insert, search, prefix matching, and autocomplete.

t = Trie() t.insert("apple") t.search("apple") True t.starts_with("app") True t.search("app") False

Source code in src/algo/trees/trie.py 
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
class Trie:
    """Prefix tree supporting insert, search, prefix matching, and autocomplete.

    >>> t = Trie()
    >>> t.insert("apple")
    >>> t.search("apple")
    True
    >>> t.starts_with("app")
    True
    >>> t.search("app")
    False
    """

    def __init__(self) -> None:
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        """Insert a word into the trie.

        >>> t = Trie()
        >>> t.insert("cat")
        >>> t.search("cat")
        True
        """
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.is_end = True

    def search(self, word: str) -> bool:
        """Return True if the exact word exists in the trie.

        >>> t = Trie()
        >>> t.insert("hello")
        >>> t.search("hello")
        True
        >>> t.search("hell")
        False
        """
        node = self._find_node(word)
        return node is not None and node.is_end

    def starts_with(self, prefix: str) -> bool:
        """Return True if any word in the trie starts with the given prefix.

        >>> t = Trie()
        >>> t.insert("hello")
        >>> t.starts_with("hel")
        True
        >>> t.starts_with("xyz")
        False
        """
        return self._find_node(prefix) is not None

    def delete(self, word: str) -> bool:
        """Remove a word from the trie, cleaning up empty nodes.

        Returns True if the word was found and deleted, False otherwise.

        >>> t = Trie()
        >>> t.insert("cat")
        >>> t.delete("cat")
        True
        >>> t.search("cat")
        False
        """
        return self._delete(self.root, word, 0)

    def autocomplete(self, prefix: str, limit: int = 10) -> list[str]:
        """Return up to `limit` words starting with the given prefix.

        >>> t = Trie()
        >>> for w in ["bar", "bat", "ball"]:
        ...     t.insert(w)
        >>> sorted(t.autocomplete("ba"))
        ['ball', 'bar', 'bat']
        """
        node = self._find_node(prefix)
        if node is None:
            return []
        results: list[str] = []
        self._collect(node, prefix, limit, results)
        return results

    def _find_node(self, prefix: str) -> TrieNode | None:
        """Traverse the trie following the prefix, returning the final node."""
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return None
            node = node.children[ch]
        return node

    def _delete(self, node: TrieNode, word: str, depth: int) -> bool:
        """Recursively delete a word and prune empty branches."""
        if depth == len(word):
            if not node.is_end:
                return False
            node.is_end = False
            return True

        ch = word[depth]
        if ch not in node.children:
            return False

        deleted = self._delete(node.children[ch], word, depth + 1)
        if deleted:
            child = node.children[ch]
            if not child.is_end and not child.children:
                del node.children[ch]
        return deleted

    def _collect(
        self,
        node: TrieNode,
        prefix: str,
        limit: int,
        results: list[str],
    ) -> None:
        """DFS to collect words under a node up to the limit."""
        if len(results) >= limit:
            return
        if node.is_end:
            results.append(prefix)
        for ch in sorted(node.children):
            if len(results) >= limit:
                return
            self._collect(node.children[ch], prefix + ch, limit, results)

autocomplete 

autocomplete(prefix: str, limit: int = 10) -> list[str]

Return up to limit words starting with the given prefix.

t = Trie() for w in ["bar", "bat", "ball"]: ... t.insert(w) sorted(t.autocomplete("ba")) ['ball', 'bar', 'bat']

Source code in src/algo/trees/trie.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def autocomplete(self, prefix: str, limit: int = 10) -> list[str]:
    """Return up to `limit` words starting with the given prefix.

    >>> t = Trie()
    >>> for w in ["bar", "bat", "ball"]:
    ...     t.insert(w)
    >>> sorted(t.autocomplete("ba"))
    ['ball', 'bar', 'bat']
    """
    node = self._find_node(prefix)
    if node is None:
        return []
    results: list[str] = []
    self._collect(node, prefix, limit, results)
    return results

delete 

delete(word: str) -> bool

Remove a word from the trie, cleaning up empty nodes.

Returns True if the word was found and deleted, False otherwise.

t = Trie() t.insert("cat") t.delete("cat") True t.search("cat") False

Source code in src/algo/trees/trie.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def delete(self, word: str) -> bool:
    """Remove a word from the trie, cleaning up empty nodes.

    Returns True if the word was found and deleted, False otherwise.

    >>> t = Trie()
    >>> t.insert("cat")
    >>> t.delete("cat")
    True
    >>> t.search("cat")
    False
    """
    return self._delete(self.root, word, 0)

insert 

insert(word: str) -> None

Insert a word into the trie.

t = Trie() t.insert("cat") t.search("cat") True

Source code in src/algo/trees/trie.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def insert(self, word: str) -> None:
    """Insert a word into the trie.

    >>> t = Trie()
    >>> t.insert("cat")
    >>> t.search("cat")
    True
    """
    node = self.root
    for ch in word:
        if ch not in node.children:
            node.children[ch] = TrieNode()
        node = node.children[ch]
    node.is_end = True

search 

search(word: str) -> bool

Return True if the exact word exists in the trie.

t = Trie() t.insert("hello") t.search("hello") True t.search("hell") False

Source code in src/algo/trees/trie.py 
64
65
66
67
68
69
70
71
72
73
74
75
def search(self, word: str) -> bool:
    """Return True if the exact word exists in the trie.

    >>> t = Trie()
    >>> t.insert("hello")
    >>> t.search("hello")
    True
    >>> t.search("hell")
    False
    """
    node = self._find_node(word)
    return node is not None and node.is_end

starts_with 

starts_with(prefix: str) -> bool

Return True if any word in the trie starts with the given prefix.

t = Trie() t.insert("hello") t.starts_with("hel") True t.starts_with("xyz") False

Source code in src/algo/trees/trie.py 
77
78
79
80
81
82
83
84
85
86
87
def starts_with(self, prefix: str) -> bool:
    """Return True if any word in the trie starts with the given prefix.

    >>> t = Trie()
    >>> t.insert("hello")
    >>> t.starts_with("hel")
    True
    >>> t.starts_with("xyz")
    False
    """
    return self._find_node(prefix) is not None
tests/trees/test_trie.py
"""Tests for the trie (prefix tree) implementation."""

from hypothesis import given, settings
from hypothesis import strategies as st

from algo.trees.trie import Trie


class TestInsertAndSearch:
    def test_basic_insert_and_search(self) -> None:
        t = Trie()
        t.insert("apple")
        assert t.search("apple") is True

    def test_search_missing_word(self) -> None:
        t = Trie()
        t.insert("apple")
        assert t.search("orange") is False

    def test_search_prefix_is_not_word(self) -> None:
        t = Trie()
        t.insert("apple")
        assert t.search("app") is False

    def test_empty_string(self) -> None:
        t = Trie()
        t.insert("")
        assert t.search("") is True

    def test_multiple_words(self) -> None:
        t = Trie()
        for word in ["cat", "car", "card", "care"]:
            t.insert(word)
        assert t.search("car") is True
        assert t.search("card") is True
        assert t.search("ca") is False


class TestStartsWith:
    def test_prefix_match(self) -> None:
        t = Trie()
        t.insert("apple")
        assert t.starts_with("app") is True

    def test_full_word_as_prefix(self) -> None:
        t = Trie()
        t.insert("apple")
        assert t.starts_with("apple") is True

    def test_empty_prefix_matches_any(self) -> None:
        t = Trie()
        t.insert("hello")
        assert t.starts_with("") is True

    def test_no_match(self) -> None:
        t = Trie()
        t.insert("apple")
        assert t.starts_with("xyz") is False

    def test_empty_trie(self) -> None:
        t = Trie()
        assert t.starts_with("a") is False


class TestDelete:
    def test_delete_existing_word(self) -> None:
        t = Trie()
        t.insert("apple")
        assert t.delete("apple") is True
        assert t.search("apple") is False

    def test_delete_nonexistent_word(self) -> None:
        t = Trie()
        t.insert("apple")
        assert t.delete("orange") is False

    def test_delete_prefix_of_another_word(self) -> None:
        t = Trie()
        t.insert("app")
        t.insert("apple")
        assert t.delete("app") is True
        assert t.search("app") is False
        assert t.search("apple") is True

    def test_delete_word_sharing_prefix(self) -> None:
        t = Trie()
        t.insert("app")
        t.insert("apple")
        assert t.delete("apple") is True
        assert t.search("apple") is False
        assert t.search("app") is True

    def test_delete_cleans_empty_nodes(self) -> None:
        t = Trie()
        t.insert("abc")
        t.delete("abc")
        assert not t.root.children


class TestAutocomplete:
    def test_returns_matching_words(self) -> None:
        t = Trie()
        for word in ["bat", "bar", "ball", "cat"]:
            t.insert(word)
        results = t.autocomplete("ba")
        assert sorted(results) == ["ball", "bar", "bat"]

    def test_respects_limit(self) -> None:
        t = Trie()
        for word in ["aa", "ab", "ac", "ad", "ae"]:
            t.insert(word)
        results = t.autocomplete("a", limit=3)
        assert len(results) == 3

    def test_empty_prefix_returns_all(self) -> None:
        t = Trie()
        words = ["cat", "car", "bat"]
        for word in words:
            t.insert(word)
        results = t.autocomplete("")
        assert sorted(results) == sorted(words)

    def test_no_matches(self) -> None:
        t = Trie()
        t.insert("apple")
        assert t.autocomplete("xyz") == []

    def test_exact_word_included(self) -> None:
        t = Trie()
        t.insert("app")
        t.insert("apple")
        results = t.autocomplete("app")
        assert sorted(results) == ["app", "apple"]


class TestTrieProperties:
    @given(
        words=st.lists(
            st.text(alphabet="abcdefghij", min_size=1, max_size=10),
            min_size=1,
            max_size=20,
        ),
    )
    @settings(max_examples=50)
    def test_inserted_words_are_searchable(self, words: list[str]) -> None:
        t = Trie()
        for word in words:
            t.insert(word)
        for word in words:
            assert t.search(word) is True

    @given(
        words=st.lists(
            st.text(alphabet="abc", min_size=1, max_size=6),
            min_size=1,
            max_size=15,
        ),
        query=st.text(alphabet="xyz", min_size=1, max_size=6),
    )
    @settings(max_examples=50)
    def test_non_inserted_words_not_found(self, words: list[str], query: str) -> None:
        t = Trie()
        for word in words:
            t.insert(word)
        assert t.search(query) is False

Implement it yourself

Run: just challenge trees trie

Then implement the functions to make all tests pass. Use just study trees for watch mode.

Reveal Solution

trie

Trie (prefix tree) for efficient string operations.

Problem

Implement a prefix tree that supports insert, search, prefix matching, delete, and autocomplete operations on a set of strings.

Approach

A tree where each node represents a character. Paths from the root to nodes marked as word-endings form the stored words. Children are stored in a dict keyed by character for O(1) branching.

When to use

Autocomplete, spell checking, IP routing, prefix matching. ASI relevance: airport code lookup (e.g., "BO" -> ["BOS", "BOI", "BOG"]), flight ID prefix search, geospatial name indexing.

Complexity

Time: O(m) per insert/search/starts_with/delete where m = word length. O(p + k) for autocomplete where p = prefix length, k = total characters in matching subtree. Space: O(N * m) where N = number of words, m = average word length.

Trie

Prefix tree supporting insert, search, prefix matching, and autocomplete.

t = Trie() t.insert("apple") t.search("apple") True t.starts_with("app") True t.search("app") False

Source code in src/algo/trees/trie.py 
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
class Trie:
    """Prefix tree supporting insert, search, prefix matching, and autocomplete.

    >>> t = Trie()
    >>> t.insert("apple")
    >>> t.search("apple")
    True
    >>> t.starts_with("app")
    True
    >>> t.search("app")
    False
    """

    def __init__(self) -> None:
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        """Insert a word into the trie.

        >>> t = Trie()
        >>> t.insert("cat")
        >>> t.search("cat")
        True
        """
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.is_end = True

    def search(self, word: str) -> bool:
        """Return True if the exact word exists in the trie.

        >>> t = Trie()
        >>> t.insert("hello")
        >>> t.search("hello")
        True
        >>> t.search("hell")
        False
        """
        node = self._find_node(word)
        return node is not None and node.is_end

    def starts_with(self, prefix: str) -> bool:
        """Return True if any word in the trie starts with the given prefix.

        >>> t = Trie()
        >>> t.insert("hello")
        >>> t.starts_with("hel")
        True
        >>> t.starts_with("xyz")
        False
        """
        return self._find_node(prefix) is not None

    def delete(self, word: str) -> bool:
        """Remove a word from the trie, cleaning up empty nodes.

        Returns True if the word was found and deleted, False otherwise.

        >>> t = Trie()
        >>> t.insert("cat")
        >>> t.delete("cat")
        True
        >>> t.search("cat")
        False
        """
        return self._delete(self.root, word, 0)

    def autocomplete(self, prefix: str, limit: int = 10) -> list[str]:
        """Return up to `limit` words starting with the given prefix.

        >>> t = Trie()
        >>> for w in ["bar", "bat", "ball"]:
        ...     t.insert(w)
        >>> sorted(t.autocomplete("ba"))
        ['ball', 'bar', 'bat']
        """
        node = self._find_node(prefix)
        if node is None:
            return []
        results: list[str] = []
        self._collect(node, prefix, limit, results)
        return results

    def _find_node(self, prefix: str) -> TrieNode | None:
        """Traverse the trie following the prefix, returning the final node."""
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return None
            node = node.children[ch]
        return node

    def _delete(self, node: TrieNode, word: str, depth: int) -> bool:
        """Recursively delete a word and prune empty branches."""
        if depth == len(word):
            if not node.is_end:
                return False
            node.is_end = False
            return True

        ch = word[depth]
        if ch not in node.children:
            return False

        deleted = self._delete(node.children[ch], word, depth + 1)
        if deleted:
            child = node.children[ch]
            if not child.is_end and not child.children:
                del node.children[ch]
        return deleted

    def _collect(
        self,
        node: TrieNode,
        prefix: str,
        limit: int,
        results: list[str],
    ) -> None:
        """DFS to collect words under a node up to the limit."""
        if len(results) >= limit:
            return
        if node.is_end:
            results.append(prefix)
        for ch in sorted(node.children):
            if len(results) >= limit:
                return
            self._collect(node.children[ch], prefix + ch, limit, results)

autocomplete 

autocomplete(prefix: str, limit: int = 10) -> list[str]

Return up to limit words starting with the given prefix.

t = Trie() for w in ["bar", "bat", "ball"]: ... t.insert(w) sorted(t.autocomplete("ba")) ['ball', 'bar', 'bat']

Source code in src/algo/trees/trie.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def autocomplete(self, prefix: str, limit: int = 10) -> list[str]:
    """Return up to `limit` words starting with the given prefix.

    >>> t = Trie()
    >>> for w in ["bar", "bat", "ball"]:
    ...     t.insert(w)
    >>> sorted(t.autocomplete("ba"))
    ['ball', 'bar', 'bat']
    """
    node = self._find_node(prefix)
    if node is None:
        return []
    results: list[str] = []
    self._collect(node, prefix, limit, results)
    return results

delete 

delete(word: str) -> bool

Remove a word from the trie, cleaning up empty nodes.

Returns True if the word was found and deleted, False otherwise.

t = Trie() t.insert("cat") t.delete("cat") True t.search("cat") False

Source code in src/algo/trees/trie.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def delete(self, word: str) -> bool:
    """Remove a word from the trie, cleaning up empty nodes.

    Returns True if the word was found and deleted, False otherwise.

    >>> t = Trie()
    >>> t.insert("cat")
    >>> t.delete("cat")
    True
    >>> t.search("cat")
    False
    """
    return self._delete(self.root, word, 0)

insert 

insert(word: str) -> None

Insert a word into the trie.

t = Trie() t.insert("cat") t.search("cat") True

Source code in src/algo/trees/trie.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def insert(self, word: str) -> None:
    """Insert a word into the trie.

    >>> t = Trie()
    >>> t.insert("cat")
    >>> t.search("cat")
    True
    """
    node = self.root
    for ch in word:
        if ch not in node.children:
            node.children[ch] = TrieNode()
        node = node.children[ch]
    node.is_end = True

search 

search(word: str) -> bool

Return True if the exact word exists in the trie.

t = Trie() t.insert("hello") t.search("hello") True t.search("hell") False

Source code in src/algo/trees/trie.py 
64
65
66
67
68
69
70
71
72
73
74
75
def search(self, word: str) -> bool:
    """Return True if the exact word exists in the trie.

    >>> t = Trie()
    >>> t.insert("hello")
    >>> t.search("hello")
    True
    >>> t.search("hell")
    False
    """
    node = self._find_node(word)
    return node is not None and node.is_end

starts_with 

starts_with(prefix: str) -> bool

Return True if any word in the trie starts with the given prefix.

t = Trie() t.insert("hello") t.starts_with("hel") True t.starts_with("xyz") False

Source code in src/algo/trees/trie.py 
77
78
79
80
81
82
83
84
85
86
87
def starts_with(self, prefix: str) -> bool:
    """Return True if any word in the trie starts with the given prefix.

    >>> t = Trie()
    >>> t.insert("hello")
    >>> t.starts_with("hel")
    True
    >>> t.starts_with("xyz")
    False
    """
    return self._find_node(prefix) is not None