LeetCode_Replace Words_TRIE.py

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

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

    def insert(self, word):
        curr = self.root
        for c in word:
            if c not in curr.children:
                curr.children[c] = TrieNode()
            curr = curr.children[c]
        curr.isWord = True

    def prefixSearch(self, word):
        curr = self.root
        for idx, c in enumerate(word):
            if c not in curr.children:
                return None
            curr = curr.children[c]
            if curr.isWord is True:
                return word[0:idx+1]
        return None

    
class Solution:
    def replaceWords(self, list_of_roots, sentence):
        """
        :type dict: List[str]
        :type sentence: str
        :rtype: str
        """
        worddict = WordDictionary()
        for root in list_of_roots:
            worddict.insert(root)
        
        result = []
        for word in sentence.split():
            is_there_prefix = worddict.prefixSearch(word)
            if is_there_prefix is None:
                result.append(word)
            else:
                result.append(is_there_prefix)
        return ' '.join(result)