From dbc2c4f581b6bc22ce431d171b87147b57618953 Mon Sep 17 00:00:00 2001 From: beam Date: Sun, 25 Oct 2026 20:12:34 +0000 Subject: [PATCH] beam: add Trie (prefix tree) with tests --- wizard-work/beam/test_trie.py | 73 ++++++++++++++++++++++ wizard-work/beam/trie.py | 112 ++++++++++++++++++++++++++++++++++ 2 files changed, 185 insertions(+) create mode 100644 wizard-work/beam/test_trie.py create mode 100644 wizard-work/beam/trie.py diff --git a/wizard-work/beam/test_trie.py b/wizard-work/beam/test_trie.py new file mode 100644 index 0000000..879d0cc --- /dev/null +++ b/wizard-work/beam/test_trie.py @@ -0,0 +1,73 @@ +import random +from trie import Trie + + +def test_basic(): + t = Trie() + for w in ["cat", "car", "card", "dog", ""]: + t.insert(w) + assert t.search("cat") + assert t.search("") + assert t.search("card") + assert not t.search("ca") + assert t.starts_with("ca") + assert not t.starts_with("z") + assert sorted(t.words_with_prefix("ca")) == ["car", "card", "cat"] + assert t.count_prefix("ca") == 3 + assert len(t) == 5 + + +def test_delete(): + t = Trie() + for w in ["car", "card", "care"]: + t.insert(w) + assert t.delete("car") + assert not t.search("car") + assert t.search("card") + assert t.search("care") + assert t.starts_with("car") + assert not t.delete("car") + # delete shared prefix leaf should prune + t.delete("card"); t.delete("care") + assert not t.starts_with("car") + + +def test_duplicate_insert(): + t = Trie() + t.insert("abc"); t.insert("abc") + assert t.search("abc") + t.delete("abc") + assert not t.search("abc") # one delete removes word entirely (count not tracked per multiplicity) + # acceptable: delete removes end flag + + +def test_random_vs_bruteforce(): + rng = random.Random(123) + words = [] + alphabet = "ab" + for _ in range(200): + w = "".join(rng.choice(alphabet) for _ in range(rng.randint(0, 5))) + words.append(w) + t = Trie() + s = set() + for w in words: + t.insert(w); s.add(w) + for w in set(words): + assert t.search(w) == (w in s) + # delete half + to_del = list(s)[: len(s) // 2] + for w in to_del: + assert t.delete(w) + s.discard(w) + for w in set(words): + assert t.search(w) == (w in s), (w, t.search(w), w in s) + # prefix enumeration matches + for p in ["", "a", "b", "aa", "ab"]: + got = set(t.words_with_prefix(p)) + exp = {w for w in s if w.startswith(p)} + assert got == exp, (p, got, exp) + + +if __name__ == "__main__": + test_basic(); test_duplicate_insert(); test_delete(); test_random_vs_bruteforce() + print("ALL OK") diff --git a/wizard-work/beam/trie.py b/wizard-work/beam/trie.py new file mode 100644 index 0000000..f1dc0dc --- /dev/null +++ b/wizard-work/beam/trie.py @@ -0,0 +1,112 @@ +"""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"))