======= (old) string.py ====== ===== Trie v1.0 ===== # 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())