Problem Statement in English

Your task is to implement a trie data structure.

Approach

A Trie allows us to store strings in a way that allows us to search for them efficiently. Prefix searches are particularly efficient. Why? Let’s see.

First let’s build a Trie. We’ll work through an example, and hopefully then, it should make sense to you.

Let’s build a Trie for the word “trie”.

    graph TD
    A(' ') --> B('t', False)
    B --> C('r', False)

    C --> D('i', False)
    D --> E('e', True)

Let’s also add the word “trial” to the same Trie.

    graph TD
    A(' ') --> B('t', False)
    B --> C('r', False)

    C --> D('i', False)
    D --> E('e', True)

    D --> F('a', False)
    F --> G('l', True)

How about “treat”?

    graph TD
    A(' ') --> B('t', False)
    B --> C('r', False)

    C --> D('i', False)
    D --> E('e', True)

    D --> F('a', False)
    F --> G('l', True)

    C --> H('e', False)
    H --> I('a', False)
    I --> J('t', True)

And lastly, let’s add the word “great” to the same Trie.

    graph TD
    A(' ') --> B('t', False)
    B --> C('r', False)

    C --> D('i', False)
    D --> E('e', True)

    D --> F('a', False)
    F --> G('l', True)

    C --> H('e', False)
    H --> I('a', False)
    I --> J('t', True)

    A --> K('g', False)
    K --> L('r', False)
    L --> M('e', False)
    M --> N('a', False)
    N --> O('t', True)

Notice how the boolean value indicates that that particular node is the end of a word. This is quite useful when we are searching for a word, but it also serves another purpose.

For example, if we added the word “team” to the Trie, we technically have “tea” as well in the Trie, but we didn’t add it. So if “tea” is searched for, we need to return False, and these boolean values help us do exactly that.

Solution in Python


class TrieNode:
    def __init__(self) -> None:
        # A dictionary to store the child tries of the node
        self.children: dict[str, TrieNode] = {}
        # A boolean to indicate if the node is the end of a word
        self.isTerminal = False


class PrefixTree:

    def __init__(self):
        # The root of the trie
        # empty, essentially
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        # Start at the root
        ptr = self.root

        # For each character in the word
        # we wish to search for
        for char in word:
            # If the character is not in the
            # children of the current Trie node
            if char not in ptr.children:
                # Add the character to the children
                # as a new Trie node
                ptr.children[char] = TrieNode()

            # Move the pointer to the child node
            # corresponding to the current character
            ptr = ptr.children[char]

        # Mark the last node as the end of a word
        ptr.isTerminal = True


    def search(self, word: str) -> bool:
        # Start at the root
        ptr = self.root

        # For each character in the word
        for char in word:
            # If the character is in the children
            # of the current Trie node
            if char in ptr.children:
                # Move the pointer to the child trie node
                ptr = ptr.children[char]
            else:
                # If the character is not in the children
                # of the current Trie node, the word is not
                # in the Trie
                return False

        # If the last node is marked as the end of a word
        # then the word is in the Trie
        return ptr.isTerminal

    def startsWith(self, prefix: str) -> bool:
        # Start at the root
        ptr = self.root

        # For each character in the word
        for char in prefix:
            # If the character is in the children
            # of the current Trie node
            if char in ptr.children:
                # Move the pointer to the child trie node
                ptr = ptr.children[char]
            # If the character is not in the children
            # of the current Trie node, the prefix is not
            else:
                # return `False`
                return False

        # if we've got this far, the prefix is in the Trie
        # return `True`
        return True

Complexity

  • Time:

    • Insert: $O(n)$
      Since we insert $n$ characters at most in the Trie.

    • Search: $O(n)$
      Since we search through $n$ characters at most in the Trie.

    • Starts With: $O(n)$
      Since we search through $n$ characters at most in the Trie.

    Where $n$ is the length of the word

  • Space: $O(n)$
    Since we store $n$ characters at most in the Trie.

And we are done.