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