Trie

trie.py

__author__ = "Shihang Zhang"
__status__ = "Prototype"


class TrieNode(object):
    """ Trie Node Class

    @attrs:
    children: a dictionary contains all children nodes
    count: an integer indicates the number of occurence of current word
    """
    def __init__(self):
        self.children = {}
        self.count = 0


class Trie(object):
    """ Trie Class

    @attrs:
    root: the root node of the trie
    """
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        """
        Inserts a word into the trie.

        @params
        prefix: string

        @return
        a bool value
        """
        current = self.root

        for c in word:
            current = current.children.setdefault(c, TrieNode())

        current.count += 1

    def search(self, word):
        """
        Returns if the word is in the trie.

        @params
        word: string

        @return
        a bool value
        """
        current = self.root

        for c in word:
            if c in current.children:
                current = current.children[c]
            else:
                return False

        if current.count > 0:
            return True
        else:
            return False

    def starts_with(self, prefix):
        """
        Returns if there is any word in the trie
        that starts with the given prefix.

        @params
        prefix: string

        @return
        a bool value
        """
        current = self.root

        for c in prefix:
            if c in current.children:
                current = current.children[c]
            else:
                return False

        return True