Source code for titan_pylib.string.trie


[docs]
 1class Trie:

[docs]
 2    class Node:
 3        def __init__(self):
 4            self.c = None
 5            self.child = {}
 6            self.count = 0
 7            self.stop_count = 0

 8
 9    def __init__(self):
10        self.root = Trie.Node()
11

[docs]
12    def add(self, s: str) -> None:
13        node = self.root
14        for dep, c in enumerate(s):
15            node.count += 1
16            if c not in node.child:
17                node.child[c] = Trie.Node()
18            node = node.child[c]
19            node.c = c
20        node.stop_count += 1

21

[docs]
22    def s_prefix(self, pref) -> int:
23        node = self.root
24        for dep, c in enumerate(pref):
25            if c not in node.child:
26                return dep
27            node = node.child[c]
28        return len(pref)

29

[docs]
30    def count(self, s: str) -> int:
31        node = self.root
32        for dep, c in enumerate(s):
33            if c not in node.child:
34                return 0
35            node = node.child[c]
36        return node.stop_count

37
38    def __contains__(self, s: str) -> bool:
39        return self.count(s) > 0
40

[docs]
41    def print(self, is_sort=False) -> None:
42        def dfs(node: Trie.Node, indent: str) -> None:
43            if len(node.child) == 0:
44                return
45            a = list(node.child.items())
46            if is_sort:
47                a.sort()  # 挿入順にするかどうか
48            for c, child in a[:-1]:
49                if child.stop_count > 0:
50                    c = "\033[32m" + c + "\033[m"
51                print(f"{indent}├── {c}")
52                dfs(child, f"{indent}|   ")
53            c, child = a[-1]
54            if child.stop_count > 0:
55                c = "\033[32m" + c + "\033[m"
56            print(f"{indent}└── {c}")
57            dfs(child, f"{indent}    ")
58
59        print("root")
60        dfs(self.root, "")