diff --git a/leetcode/stuff.py b/leetcode/stuff.py index e4fa621..7bc537a 100644 --- a/leetcode/stuff.py +++ b/leetcode/stuff.py @@ -473,6 +473,8 @@ class Trie: cur_t.kids.setdefault(prefix, next_val) cur_t = cur_t.kids[prefix] + cur_t.kids["__leaf__"] = TrieNode("__leaf__", {}) + def search(self, word: str) -> bool: return self._has(word, prefix_ok=False) @@ -484,12 +486,14 @@ class Trie: cur_t = self._t value = cur_t.value + is_leaf: bool = False 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 and value == word + return (cur_t is None or is_leaf) and value == word