trie tree demo

trie_tree.py

class TrieNode(object):

    def __init__(self):
        self.children = dict()
        self.is_word = False


class TrieTree(object):

    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        node = self.root
        for letter in word:
            child = node.children.get(letter)
            if child is None:
                child = TrieNode()
                node.children[letter] = child
            node = child
        node.is_word = True

    def search(self, word):
        ans = ''
        node = self.root
        for letter in word:
            child = node.children.get(letter)
            if child is None:
                return False
            ans += letter
            if child.is_word:
                print('contain word: %s' % ans)
            node = child
        return True


if __name__ == '__main__':
    tree = TrieTree()   
    words = ['a', 'the', 'there', 'their', 'answser', 'any', 'bye', 'by']
    for w in words:
        tree.insert(w)
    print(tree.search('there'))