Fully un-busted yet slow Trie

leetcode/stuff.py

@@ -445,3 +445,56 @@ class RandomizedSet:
 
     def getRandom(self) -> int:
         return random.choice(self._l)
+
+
+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._root_node = TrieNode("", {})
+
+    def insert(self, word: str) -> None:
+        if len(word) == 0:
+            return
+
+        current_node = self._root_node
+
+        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]
+
+        leaf = TrieNode.leaf()
+        current_node.kids[leaf.value] = leaf
+
+    def search(self, word: str) -> bool:
+        return self._has(word, prefix_ok=False)
+
+    def startsWith(self, prefix: str) -> bool:
+        return self._has(prefix, prefix_ok=True)
+
+    def _has(self, word: str, prefix_ok: bool) -> bool:
+        if len(word) == 0:
+            return True
+
+        reverse_path = [word[: i + 1] for i in range(len(word))][::-1]
+        current_node = self._root_node
+
+        while reverse_path and current_node is not None:
+            current_node = current_node.kids.get(reverse_path.pop())
+
+        return (
+            current_node is not None
+            and (current_node.is_leaf or prefix_ok)
+            and current_node.value == word
+        )

leetcode/test_stuff.py

@@ -388,3 +388,59 @@ def test_randomized_set(cls: type[stuff.RandomizedSet] | type[stuff.SlowRandomiz
     assert inst.insert(2) is True
     assert inst.remove(1) is True
     assert inst.getRandom() == 2
+
+
+def test_trie_single_letter():
+    trie = stuff.Trie()
+
+    assert trie.insert("a") is None
+    assert trie.search("a") is True
+    assert trie.startsWith("a") is True
+
+
+def test_trie_prefix_leaf():
+    trie = stuff.Trie()
+
+    assert trie.insert("apple") is None
+    assert trie.search("apple") is True
+    assert trie.search("app") is False
+    assert trie.startsWith("app") is True
+    assert trie.insert("app") is None
+    assert trie.search("app") is True
+
+
+def test_trie_two_letter():
+    trie = stuff.Trie()
+
+    assert trie.insert("ab") is None
+    assert trie.search("a") is False
+    assert trie.startsWith("a") is True
+
+
+def test_trie_busy():
+    trie = stuff.Trie()
+
+    assert trie.insert("app") is None
+    assert trie.insert("apple") is None
+    assert trie.insert("beer") is None
+    assert trie.insert("add") is None
+    assert trie.insert("jam") is None
+    assert trie.insert("rental") is None
+    assert trie.search("apps") is False
+    assert trie.search("app") is True
+    assert trie.search("ad") is False
+    assert trie.search("applepie") is False
+    assert trie.search("rest") is False
+    assert trie.search("jan") is False
+    assert trie.search("rent") is False
+    assert trie.search("beer") is True
+    assert trie.search("jam") is True
+    assert trie.startsWith("apps") is False
+    assert trie.startsWith("app") is True
+    assert trie.startsWith("ad") is True
+    assert trie.startsWith("applepie") is False
+    assert trie.startsWith("rest") is False
+    assert trie.startsWith("jan") is False
+    assert trie.startsWith("rent") is True
+    assert trie.startsWith("beer") is True
+    assert trie.startsWith("jam") is True