Fully un-busted yet slow Trie

leetcode/stuff.py

@@ -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]
+        if len(word) == 0:
-        cur_t = self._t
+            return
 
-        for i, prefix in enumerate(prefixes):
+        current_node = self._root_node
-            next_val: TrieNode | None = None
-            if i + 1 < len(prefixes) - 1:
-                next_val = TrieNode(prefixes[i + 1], {})
-            else:
-                next_val = TrieNode(word, {})
 
-            if cur_t.kids is None:
+        for prefix in [word[: i + 1] for i in range(len(word))]:
-                cur_t.kids = {}
+            current_node.kids.setdefault(prefix, TrieNode(prefix, {}))
+            current_node = current_node.kids[prefix]
 
-            cur_t.kids.setdefault(prefix, next_val)
+        leaf = TrieNode.leaf()
-            cur_t = cur_t.kids[prefix]
+        current_node.kids[leaf.value] = leaf
-
-        cur_t.kids["__leaf__"] = TrieNode("__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:
+        if len(word) == 0:
+            return True
+
         reverse_path = [word[: i + 1] for i in range(len(word))][::-1]
+        current_node = self._root_node
 
-        cur_t = self._t
+        while reverse_path and current_node is not None:
-        value = cur_t.value
+            current_node = current_node.kids.get(reverse_path.pop())
-        is_leaf: bool = False
 
-        while reverse_path and cur_t is not None:
+        return (
-            value = cur_t.value
+            current_node is not None
-            is_leaf = "__leaf__" in cur_t.kids
+            and (current_node.is_leaf or prefix_ok)
-            cur_t = cur_t.kids.get(reverse_path.pop())
+            and current_node.value == word
-
+        )
-            if prefix_ok and cur_t is not None and value == word:
-                return True
-
-        return (cur_t is None or is_leaf) and value == word

leetcode/test_stuff.py

@@ -390,7 +390,15 @@ def test_randomized_set(cls: type[stuff.RandomizedSet] | type[stuff.SlowRandomiz
     assert inst.getRandom() == 2
 
 
-def test_trie():
+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
@@ -399,3 +407,40 @@ def test_trie():
     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