Skip to content

Instantly share code, notes, and snippets.

@devbruce

devbruce/trie.py

Last active June 9, 2024 06:38
Download ZIP

Revisions

  1. devbruce revised this gist Jun 9, 2024. 1 changed file with 2 additions and 2 deletions.
    4 changes: 2 additions & 2 deletions trie.py
    View file
    Original file line number Diff line number Diff line change
    @@ -39,8 +39,8 @@ def startswith(self, keyword: str) -> list[str]:
    def _dfs(node: Node, prefix: str):
    if node.flag is True:
    candidates.append(prefix)
    for char, node in node.childrens.items():
    _dfs(node, prefix + char)
    for char, child_node in node.childrens.items():
    _dfs(child_node, prefix + char)

    _dfs(cur_node, keyword)
    return candidates
  2. devbruce created this gist Jun 8, 2024.
    81 changes: 81 additions & 0 deletions trie.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,81 @@
    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, node in node.childrens.items():
    _dfs(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']