사용자 도구

사이트 도구


ps:teflib:string_old

(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())

토론

댓글을 입력하세요:
L B F K X
 
ps/teflib/string_old.txt · 마지막으로 수정됨: 2021/01/06 04:05 저자 teferi