113 lines
3.3 KiB
Python
113 lines
3.3 KiB
Python
"""Prefix trie with insert/search/startsWith/delete and prefix enumeration.
|
|
|
|
Correctness notes:
|
|
- delete only removes nodes that are not shared by other words; marks end-of-word
|
|
flags rather than blindly removing.
|
|
- search is exact word match; starts_with is prefix match.
|
|
- words_with_prefix returns sorted list (deterministic).
|
|
- Empty string is a valid word if inserted.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
|
|
class _Node:
|
|
__slots__ = ("children", "is_end", "count")
|
|
|
|
def __init__(self):
|
|
self.children: dict[str, "_Node"] = {}
|
|
self.is_end: bool = False
|
|
self.count: int = 0 # number of words passing through this node
|
|
|
|
|
|
class Trie:
|
|
def __init__(self):
|
|
self.root = _Node()
|
|
|
|
def insert(self, word: str) -> None:
|
|
node = self.root
|
|
node.count += 1
|
|
for ch in word:
|
|
node = node.children.setdefault(ch, _Node())
|
|
node.count += 1
|
|
node.is_end = True
|
|
|
|
def search(self, word: str) -> bool:
|
|
node = self._walk(word)
|
|
return node is not None and node.is_end
|
|
|
|
def starts_with(self, prefix: str) -> bool:
|
|
return self._walk(prefix) is not None
|
|
|
|
def _walk(self, s: str):
|
|
node = self.root
|
|
for ch in s:
|
|
if ch not in node.children:
|
|
return None
|
|
node = node.children[ch]
|
|
return node
|
|
|
|
def delete(self, word: str) -> bool:
|
|
# returns True if word was present and removed
|
|
node = self._walk(word)
|
|
if node is None or not node.is_end:
|
|
return False
|
|
# decrement counts along path; prune empty branches
|
|
self._remove(self.root, word, 0)
|
|
return True
|
|
|
|
def _remove(self, node: _Node, word: str, depth: int) -> bool:
|
|
if depth == len(word):
|
|
if not node.is_end:
|
|
return False
|
|
node.is_end = False
|
|
return len(node.children) == 0
|
|
ch = word[depth]
|
|
child = node.children.get(ch)
|
|
if child is None:
|
|
return False
|
|
should_delete_child = self._remove(child, word, depth + 1)
|
|
if should_delete_child:
|
|
del node.children[ch]
|
|
return (not node.is_end) and len(node.children) == 0
|
|
return False
|
|
|
|
def words_with_prefix(self, prefix: str) -> list[str]:
|
|
node = self._walk(prefix)
|
|
if node is None:
|
|
return []
|
|
out: list[str] = []
|
|
self._collect(node, prefix, out)
|
|
return sorted(out)
|
|
|
|
def _collect(self, node: _Node, prefix: str, out: list[str]) -> None:
|
|
if node.is_end:
|
|
out.append(prefix)
|
|
for ch in sorted(node.children):
|
|
self._collect(node.children[ch], prefix + ch, out)
|
|
|
|
def count_prefix(self, prefix: str) -> int:
|
|
node = self._walk(prefix)
|
|
return node.count if node is not None else 0
|
|
|
|
def __len__(self) -> int:
|
|
# number of distinct words
|
|
out = [0]
|
|
|
|
def dfs(n):
|
|
if n.is_end:
|
|
out[0] += 1
|
|
for c in n.children.values():
|
|
dfs(c)
|
|
|
|
dfs(self.root)
|
|
return out[0]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
t = Trie()
|
|
for w in ["cat", "car", "card", "dog"]:
|
|
t.insert(w)
|
|
print(t.search("car"), t.starts_with("ca"), t.words_with_prefix("ca"))
|
|
t.delete("car")
|
|
print(t.search("car"), t.search("card"))
|