Trie Data Structure

trie.py

class TrieNode(object):

    def __init__(self):
        self.children = {}
        self.is_word = False    # mark the end of a word


class Trie(object):

    def __init__(self):
        # initialize the root node
        self.root = TrieNode()

    def add(self, word):
        # add a new word to the trie
        current = self.root
        for c in word:
            if c not in current.children:
                # create new node if the character is not in children
                current.children[c] = TrieNode()
            current = current.children[c]
        # mark the end of word
        current.is_word = True

    def search(self, word):
        # return True if the word is in the trie
        current = self.root
        for c in word:
            if c not in current.children:
                return False
            current = current.children[c]
        return current.is_word

    def search_prefix(self, prefix):
        # return True if the prefix is in the trie
        current = self.root
        for c in prefix:
            if c not in current.children:
                return False
            current = current.children[c]
        # we don't need a completed word for prefix
        return True