trie.py

from typing import Self


class Node:
    def __init__(self, char: str | None = None):
        self.char = char
        self.flag: bool = False
        self.childrens: dict[str, Self] = {}


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

    def insert(self, keyword: str):
        cur_node = self.head
        for char in keyword:
            if char not in cur_node.childrens:
                cur_node.childrens[char] = Node(char)
            cur_node = cur_node.childrens[char]
        cur_node.flag = True

    def isin(self, keyword: str) -> bool:
        cur_node = self.head
        for char in keyword:
            if char not in cur_node.childrens:
                return False
            cur_node = cur_node.childrens[char]
        return cur_node.flag

    def startswith(self, keyword: str) -> list[str]:
        candidates: list[str] = []
        cur_node = self.head
        for char in keyword:
            if char not in cur_node.childrens:
                return []
            cur_node = cur_node.childrens[char]

        def _dfs(node: Node, prefix: str):
            if node.flag is True:
                candidates.append(prefix)
            for char, child_node in node.childrens.items():
                _dfs(child_node, prefix + char)

        _dfs(cur_node, keyword)
        return candidates



if __name__ == "__main__":
    trie = Trie()
    # insert
    trie.insert(keyword="abc")
    trie.insert(keyword="abcd")
    trie.insert(keyword="abcde")
    trie.insert(keyword="abcdef")

    trie.insert(keyword="b")
    trie.insert(keyword="bc")
    trie.insert(keyword="bcd")
    trie.insert(keyword="bcde")

    trie.insert(keyword="1234")
    trie.insert(keyword="333")

    # isin
    print(trie.isin("abc"))  # True
    print(trie.isin("abcdef"))  # True
    print(trie.isin("bc"))  # True
    print(trie.isin("bcde"))  # True

    print(trie.isin("zz"))  # False
    print(trie.isin("asd"))  # False
    print(trie.isin("kkkk"))  # False
    print(trie.isin("pppp"))  # False

    # startswith
    print(trie.startswith("a"))  # ['abc', 'abcd', 'abcde', 'abcdef']
    print(trie.startswith("abcde"))  # ['abcde', 'abcdef']
    print(trie.startswith("bc"))  # ['bc', 'bcd', 'bcde']
    print(trie.startswith("1"))  # ['1234']