trie.py (1341B)
1 class TreeNode: 2 def __init__(self): 3 self.is_string = False 4 self.chars = {} 5 6 7 class Trie: 8 9 def __init__(self): 10 self.root = TreeNode() 11 12 def insert(self, word: str) -> None: 13 current = self.root 14 15 for char in word: 16 17 if not char in current.chars: 18 new_node = TreeNode() 19 current.chars[char] = new_node 20 current = new_node 21 else: 22 current = current.chars[char] 23 24 current.is_string = True 25 26 def search(self, word: str) -> bool: 27 28 current = self.root 29 30 for char in word: 31 if char in current.chars: 32 current = current.chars[char] 33 else: 34 return False 35 36 return current.is_string 37 38 def startsWith(self, prefix: str) -> bool: 39 current = self.root 40 41 for char in prefix: 42 if char in current.chars: 43 current = current.chars[char] 44 else: 45 return False 46 return True 47 48 49 # Your Trie object will be instantiated and called as such: 50 # obj = Trie() 51 # obj.insert(word) 52 # param_2 = obj.search(word) 53 # param_3 = obj.startsWith(prefix) 54 55 # idea: 56 # route through the tree based on the prefix of the string. 57 # if a string exists, we then mark the node with a 'is_string' bool 58 #