Skip to main content

208. Implement Trie (Prefix Tree) - LeetCode

from typing import Optional


class Node:
def __init__(self, is_end: bool = False):
self.is_end = is_end
self.children: dict[str, Node] = {}

def __repr__(self):
return f'''{'! ' if self.is_end else ''}{self.children}'''


class Trie:
def __init__(self):
self.root = Node()

def insert(self, word: str) -> None:
node = self.root
for c in word:
node.children[c] = node.children.get(c, Node())
node = node.children[c]
node.is_end = True

def _last_node_same_with(self, word: str) -> Optional[Node]:
node = self.root
for c in word:
if c not in node.children:
return
node = node.children[c]
return node

def search(self, word: str) -> bool:
node = self._last_node_same_with(word)
return bool(node) and node.is_end

def startsWith(self, word: str) -> bool:
node = self._last_node_same_with(word)
return bool(node)