Trie (Prefix Tree)

A Trie (pronounced as try), also known as Prefix Tree, is a Tree data structure that is used to efficiently store and locate strings with a predefined set.

As the same suggest, we can store prefixes to avoid to store common strings more than once.

For example, suppose that we have the following strings:

strings = [
    "tab",
    "table",
    "tank",
    "tart",
    "task",
]

We can notice that all of them share the same prefix ta. So, we could use a tree to store these single characters and avoid to duplicate them. Lastly, we can have a flag to mark a complete word (green node). In this way, we can distinguish between a normal character and a character that forms an actual string.

Implementation

You can implement this data structure by using a n-ary tree.

from typing import Dict, Optional


class TrieNode:
    def __init__(self, is_word: Optional[bool] = False) -> None:
        self.children: Dict[str, "TrieNode"] = {}
        self.is_word = is_word


class Trie:
    def __init__(self) -> None:
        self.head = TrieNode()

    def insert(self, word: str) -> None:
        p = self.head

        for char in word:
            if char not in p.children:
                p.children[char] = TrieNode()

            p = p.children[char]

        p.is_word = True

    def search(self, word: str) -> bool:
        return self._contains(word=word, match_word=True)

    def startsWith(self, prefix: str) -> bool:
        return self._contains(word=prefix, match_word=False)

    def _contains(self, word: str, match_word: bool) -> bool:
        p = self.head

        for char in word:
            if char not in p.children:
                return False

            p = p.children[char]

        return p.is_word if match_word else True

Operations

Let's check the complexity for the three Trie operations:

• insert:
  • Time: $O(n)$.
  • Space: $O(n)$.
• search:
  • Time: $O(n)$.
  • Space: $O(1)$.
• startsWith:
  • Time: $O(p)$.
  • Space: $O(1)$.

Where $n$ is the size of the word and $p$ is the size of the prefix.