beam: add Trie (prefix tree) with tests

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")

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"))