# N Trie
# I {"version": "1.01", "typing": ["Generator", "Generic", "Iterable", "TypeVar"]}
_CharLike = TypeVar('_CharLike')
_StrLike = Iterable[_CharLike]
class TrieNode():
def __init__(self):
self.children = {}
self.child_count = 0
self.word_count = 0
self.word = None
class Trie(Generic[_CharLike]):
def __init__(self, words: Iterable[_StrLike] = ()):
self._root = TrieNode()
for word in words:
self.add(word)
def __contains__(self, word: _StrLike):
try:
return self[word].word_count > 0
except KeyError:
return False
def __getitem__(self, prefix: _StrLike):
node = self._root
for ch in prefix:
node = node.children[ch]
return node
def add(self, word: _StrLike) -> TrieNode:
"""Adds the given word to trie and returns a node for it."""
node = self._root
for ch in word:
node.child_count += 1
node = node.children.setdefault(ch, TrieNode())
node.word_count += 1
node.word = word
return node
def prefixes(self, word: _StrLike) -> Generator[TrieNode, None, None]:
"""Yields all nodes for prefixes of the given word."""
node = self._root
return ((node := node.children[ch]) for ch in word)
def words(self, prefix: _StrLike = '') -> Generator[TrieNode, None, None]:
"""Yields all word nodes that start with the given prefix."""
try:
stack = [self[prefix]]
except KeyError:
return
while stack:
node = stack.pop()
if node.word_count > 0:
yield node
stack.extend(node.children.values())