|
|
|
@ -451,29 +451,31 @@ class TrieNode(typing.NamedTuple):
|
|
|
|
|
value: str
|
|
|
|
|
kids: dict[str, "TrieNode"]
|
|
|
|
|
|
|
|
|
|
@property
|
|
|
|
|
def is_leaf(self) -> bool:
|
|
|
|
|
return "__self__" in self.kids
|
|
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
|
def leaf(cls) -> "TrieNode":
|
|
|
|
|
return cls("__self__", {})
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Trie:
|
|
|
|
|
def __init__(self):
|
|
|
|
|
self._t = TrieNode("", {})
|
|
|
|
|
self._root_node = TrieNode("", {})
|
|
|
|
|
|
|
|
|
|
def insert(self, word: str) -> None:
|
|
|
|
|
prefixes = [word[: i + 1] for i in range(len(word))][:-1]
|
|
|
|
|
cur_t = self._t
|
|
|
|
|
|
|
|
|
|
for i, prefix in enumerate(prefixes):
|
|
|
|
|
next_val: TrieNode | None = None
|
|
|
|
|
if i + 1 < len(prefixes) - 1:
|
|
|
|
|
next_val = TrieNode(prefixes[i + 1], {})
|
|
|
|
|
else:
|
|
|
|
|
next_val = TrieNode(word, {})
|
|
|
|
|
if len(word) == 0:
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
if cur_t.kids is None:
|
|
|
|
|
cur_t.kids = {}
|
|
|
|
|
current_node = self._root_node
|
|
|
|
|
|
|
|
|
|
cur_t.kids.setdefault(prefix, next_val)
|
|
|
|
|
cur_t = cur_t.kids[prefix]
|
|
|
|
|
for prefix in [word[: i + 1] for i in range(len(word))]:
|
|
|
|
|
current_node.kids.setdefault(prefix, TrieNode(prefix, {}))
|
|
|
|
|
current_node = current_node.kids[prefix]
|
|
|
|
|
|
|
|
|
|
cur_t.kids["__leaf__"] = TrieNode("__leaf__", {})
|
|
|
|
|
leaf = TrieNode.leaf()
|
|
|
|
|
current_node.kids[leaf.value] = leaf
|
|
|
|
|
|
|
|
|
|
def search(self, word: str) -> bool:
|
|
|
|
|
return self._has(word, prefix_ok=False)
|
|
|
|
@ -482,18 +484,17 @@ class Trie:
|
|
|
|
|
return self._has(prefix, prefix_ok=True)
|
|
|
|
|
|
|
|
|
|
def _has(self, word: str, prefix_ok: bool) -> bool:
|
|
|
|
|
reverse_path = [word[: i + 1] for i in range(len(word))][::-1]
|
|
|
|
|
|
|
|
|
|
cur_t = self._t
|
|
|
|
|
value = cur_t.value
|
|
|
|
|
is_leaf: bool = False
|
|
|
|
|
if len(word) == 0:
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
while reverse_path and cur_t is not None:
|
|
|
|
|
value = cur_t.value
|
|
|
|
|
is_leaf = "__leaf__" in cur_t.kids
|
|
|
|
|
cur_t = cur_t.kids.get(reverse_path.pop())
|
|
|
|
|
reverse_path = [word[: i + 1] for i in range(len(word))][::-1]
|
|
|
|
|
current_node = self._root_node
|
|
|
|
|
|
|
|
|
|
if prefix_ok and cur_t is not None and value == word:
|
|
|
|
|
return True
|
|
|
|
|
while reverse_path and current_node is not None:
|
|
|
|
|
current_node = current_node.kids.get(reverse_path.pop())
|
|
|
|
|
|
|
|
|
|
return (cur_t is None or is_leaf) and value == word
|
|
|
|
|
return (
|
|
|
|
|
current_node is not None
|
|
|
|
|
and (current_node.is_leaf or prefix_ok)
|
|
|
|
|
and current_node.value == word
|
|
|
|
|
)
|
|
|
|
|