August - Day 5 - Add and Search Word

Solution.py

class WordDictionary:

  def __init__(self):
    """
    Initialize your data structure here.
    """
    self.root = TrieNode()


  def addWord(self, word: str) -> None:
    """
    Adds a word into the data structure.
    """
    root = self.root
    
    for c in word:
      if not root.children[ord(c) - ord('a')]:
        root.children[ord(c) - ord('a')] = TrieNode()
        
      root = root.children[ord(c) - ord('a')]
      
    root.finishes = True


  def search(self, word: str) -> bool:
    """
    Returns if the word is in the data structure. A word could contain the dot character '.' to represent any one letter.
    """
    return self.searchInternal(self.root, word, 0)
  
  def searchInternal(self, root: 'TrieNode', word: str, i: int) -> bool:
    if not root:
      return False
    
    if i == len(word) and root.finishes:
      return True
    
    if i == len(word):
      return False
    
    if word[i] == '.':
      for child in root.children:
        if self.searchInternal(child, word, i + 1):
          return True
      return False
    
    return self.searchInternal(root.children[ord(word[i]) - ord('a')], word, i + 1)
    
    
class TrieNode:
  def __init__(self):
    self.children = [None] * 26
    self.finishes = False



# Your WordDictionary object will be instantiated and called as such:
# obj = WordDictionary()
# obj.addWord(word)
# param_2 = obj.search(word)