Build a Trie in Python

Problem

In computer science, a trie (pronounced "try"), also called digital tree, radix tree or prefix tree, is a kind of search tree—an ordered tree data structure used to store a dynamic set or associative array where the keys are usually strings. Build a trie.

Input

• A class Trie, with an empty constructor and three empty methods
  • insert(word)
  • search(word)
  • startsWith(prefix):
• Assume that all word inputs are lowercase alphabets

Approach

The key to building a trie is to define a class that has a container which can house up to 26 keys (1 for each lowercase alphabet). This can be done by creating an array or hash map. Creating a static array, with a fixed size of 26 is the most optimal choice.

Each index in this array should contain another Trie object, which will contain an array of its own for its own children. A Trie is a lot like any computer science tree -- a subtree of a trie should be a trie itself.

When inserting a word, we want to iterate through each character in the word. For each character, we should create a new Trie for that character and place it in the current trie's array. If this character already exists in the array, we should keep that trie, and re-use it, rather than replace it. Replacing that trie with an entirely new node will result in lost data.

One good question is how to implement the iterating logic. Recursive code is viable for a simple implementation and it is easier to write, since tries are recursive data structures in nature. Iterative code however, is faster in the long-run albeit it being trickier to write.

Solution

This is an iterative solution of the above approach. ord is used to retrieve the ordinal number of the ASCII characters and it is subtracted with the offset to get the indices 0-25, with 0 = 'a', 1 = 'b', and so on.

ORD_OFFSET = 97  
NUM_ALPHABETS = 26

class Trie:  
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.alpha = [None] * NUM_ALPHABETS
        self.is_word = False


    def insert(self, word):
        """
        Inserts a word into the trie.
        :type word: str
        :rtype: void
        """
        letter_trie = self
        for i in range(0, len(word)):
            letter = word[i]

            # if letter already exists in this trie, re-use it
            if letter_trie.alpha[ord(letter) - ORD_OFFSET]:
                letter_trie = letter_trie.alpha[ord(letter) - ORD_OFFSET]
                continue

            new_trie = Trie()
            letter_trie.alpha[ord(letter) - ORD_OFFSET] = new_trie
            letter_trie = new_trie

        letter_trie.is_word = True


    def search(self, word, is_prefix = False):
        """
        Returns if the word is in the trie.
        :type word: str
        :rtype: bool
        """
        letter_trie = self
        for i in range(0, len(word)):
            letter = word[i]
            stored_letter = letter_trie.alpha[ord(letter) - ORD_OFFSET]
            if stored_letter is None:
                return False
            letter_trie = stored_letter

        if not is_prefix and not letter_trie.is_word:
            return False

        return True


    def startsWith(self, prefix):
        """
        Returns if there is any word in the trie that starts with the given prefix.
        :type prefix: str
        :rtype: bool
        """
        return self.search(prefix, is_prefix = True)



# Your Trie object will be instantiated and called as such:
# obj = Trie()
# obj.insert(word)
# param_2 = obj.search(word)
# param_3 = obj.startsWith(prefix)