Add prefix matching support and update test cases

Author: staging-devin-ai-integration[bot]

trie/test_trie.py

@@ -158,9 +158,23 @@ def test_prefix_matching(self):
         words = ["prefix", "preface", "prepare", "prevent"]
         for word in words:
             self.trie.insert(word)
-        self.assertTrue(all(self.trie.search(word[:i]) for word in words for i in range(3, len(word) + 1)))
+
+        # Test exact word matches
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+
+        # Test prefix matches
+        for word in words:
+            for i in range(3, len(word)):
+                self.assertTrue(self.trie.search(word[:i], is_prefix=True))
+
+        # Test that "pre" is found as a prefix but not as a complete word
+        self.assertTrue(self.trie.search("pre", is_prefix=True))
         self.assertFalse(self.trie.search("pre"))
 
+        # Test non-existent prefix
+        self.assertFalse(self.trie.search("pra", is_prefix=True))
+
     def test_edge_cases(self):
         self.trie.insert("a")
         self.assertTrue(self.trie.search("a"))

trie/trie.py

@@ -19,18 +19,18 @@ def insert(self, word):
             node = node.children[char]
         node.is_end_of_word = True
 
-    def search(self, word):
+    def search(self, word, is_prefix=False):
         """
-        Search for a word in the trie.
-        Returns True if the word is found, False otherwise.
+        Search for a word or prefix in the trie.
+        Returns True if the word/prefix is found, False otherwise.
         Time complexity: O(m), where m is the length of the word.
         """
         node = self.root
         for char in word:
             if char not in node.children:
                 return False
             node = node.children[char]
-        return node.is_end_of_word
+        return is_prefix or node.is_end_of_word
 
     def delete(self, word):
         """