Implement Trie (Prefix Tree)

trie.py

"""
Implement a trie with insert, search, and startsWith methods.
- You may assume that all inputs are consist of lowercase letters a-z.
- All inputs are guaranteed to be non-empty strings.

Example:
trie = Trie()
trie.insert("apple")
trie.search("apple")   # returns True
trie.search("app")     # returns False
trie.startsWith("app") # returns True
trie.insert("app")   
trie.search("app")     # returns True
"""

class Trie:
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.children = {}
        self.marker = '$'

    def insert(self, word):
        """
        Inserts a word into the trie.
        :type word: str
        :rtype: void
        """
        cur = self.children
        for c in word:
            if c not in cur:
                cur[c] = {}
            cur = cur[c]
        cur[self.marker] = True
    
    def __search__(self, word):
        cur = self.children
        for c in word:
            if c not in cur:
                return False
            cur = cur[c]
        return cur
        
    def search(self, word):
        """
        Returns if the word is in the trie.
        :type word: str
        :rtype: bool
        """
        found = self.__search__(word)
        return found and len(found) > 0 and self.marker in found

    def startsWith(self, prefix):
        """
        Returns if there is any word in the trie that starts with the given prefix.
        :type prefix: str
        :rtype: bool
        """
        if len(prefix) == 0: return True
        found = self.__search__(prefix)
        return found and len(found) > 0