Lintcode 442. Implement Trie (Prefix Tree) (Medium) (Python)

Implement Trie

Description:

Implement a trie with insert, search, and startsWith methods.

Example
insert(“lintcode”)
search(“code”)

false
startsWith(“lint”)
true
startsWith(“linterror”)
false
insert(“linterror”)
search(“lintcode)
true
startsWith(“linterror”)
true
Notice
You may assume that all inputs are consist of lowercase letters a-z.

Code:

class TrieNode:
    def __init__(self):
        self.child = {}
        self.isWord = False

class Trie:

    def __init__(self):
        # do intialization if necessary
        self.root = TrieNode()
    """
    @param: word: a word
    @return: nothing
    """
    def insert(self, word):
        # write your code here
        cur = self.root
        for i in word:
            if i not in cur.child:
                newNode = TrieNode()
                cur.child[i] = newNode
                cur = newNode
            else:
                cur = cur.child[i]
        cur.isWord = True
    """
    @param: word: A string
    @return: if the word is in the trie.
    """
    def search(self, word):
        # write your code here
        cur = self.root
        for i in word:
            if i not in cur.child:
                return False
            cur = cur.child[i]
        return cur.isWord

    """
    @param: prefix: A string
    @return: if there is any word in the trie that starts with the given prefix.
    """
    def startsWith(self, prefix):
        # write your code here
        cur = self.root
        for i in prefix:
            print(i)
            if i not in cur.child:
                return False
            cur = cur.child[i]
        return True