Busted trie impl

leetcode/stuff.py

@@ -445,3 +445,51 @@ class RandomizedSet:
 
     def getRandom(self) -> int:
         return random.choice(self._l)
+
+
+class TrieNode(typing.NamedTuple):
+    value: str
+    kids: dict[str, "TrieNode"]
+
+
+class Trie:
+    def __init__(self):
+        self._t = TrieNode("", {})
+
+    def insert(self, word: str) -> None:
+        prefixes = [word[: i + 1] for i in range(len(word))][:-1]
+        cur_t = self._t
+
+        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, {})
+
+            if cur_t.kids is None:
+                cur_t.kids = {}
+
+            cur_t.kids.setdefault(prefix, next_val)
+            cur_t = cur_t.kids[prefix]
+
+    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:
+        reverse_path = [word[: i + 1] for i in range(len(word))][::-1]
+
+        cur_t = self._t
+        value = cur_t.value
+
+        while reverse_path and cur_t is not None:
+            value = cur_t.value
+            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 and value == word

leetcode/test_stuff.py

@@ -388,3 +388,14 @@ 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():
+    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