Implement new Trie data structure with insert, search, and starts_wit…

…h methods

tree/trie/new_trie.py

@@ -0,0 +1,80 @@
+"""
+Implement a new Trie data structure with insert, search, and starts_with methods.
+
+This implementation uses a TrieNode class to represent each node in the Trie.
+Each TrieNode contains a dictionary of children nodes and a flag to indicate
+if it represents the end of a word.
+
+Note:
+    - All inputs are assumed to consist of lowercase letters a-z.
+    - This implementation aims to be more memory-efficient by using a regular
+      dictionary instead of defaultdict for children nodes.
+"""
+
+
+class TrieNode:
+    def __init__(self):
+        self.children = {}
+        self.is_end_of_word = False
+
+
+class Trie:
+    def __init__(self):
+        self.root = TrieNode()
+
+    def insert(self, word: str) -> None:
+        """
+        Inserts a word into the trie.
+
+        Args:
+            word (str): The word to be inserted.
+        """
+        node = self.root
+        for char in word:
+            if char not in node.children:
+                node.children[char] = TrieNode()
+            node = node.children[char]
+        node.is_end_of_word = True
+
+    def search(self, word: str) -> bool:
+        """
+        Returns True if the word is in the trie, False otherwise.
+
+        Args:
+            word (str): The word to search for.
+
+        Returns:
+            bool: True if the word is found, False otherwise.
+        """
+        node = self._find_node(word)
+        return node is not None and node.is_end_of_word
+
+    def starts_with(self, prefix: str) -> bool:
+        """
+        Returns True if there is any word in the trie that starts with the given prefix.
+
+        Args:
+            prefix (str): The prefix to search for.
+
+        Returns:
+            bool: True if any word starts with the given prefix, False otherwise.
+        """
+        return self._find_node(prefix) is not None
+
+    def _find_node(self, prefix: str) -> TrieNode:
+        """
+        Helper method to find the node corresponding to the given prefix.
+
+        Args:
+            prefix (str): The prefix to search for.
+
+        Returns:
+            TrieNode: The node corresponding to the last character of the prefix,
+                      or None if the prefix is not found.
+        """
+        node = self.root
+        for char in prefix:
+            if char not in node.children:
+                return None
+            node = node.children[char]
+        return node

tree/trie/test_new_trie.py

@@ -0,0 +1,38 @@
+import unittest
+from new_trie import Trie
+
+class TestTrie(unittest.TestCase):
+    def setUp(self):
+        self.trie = Trie()
+
+    def test_insert_and_search(self):
+        self.trie.insert("apple")
+        self.assertTrue(self.trie.search("apple"))
+        self.assertFalse(self.trie.search("app"))
+        self.assertFalse(self.trie.search("apples"))
+
+    def test_starts_with(self):
+        self.trie.insert("apple")
+        self.assertTrue(self.trie.starts_with("app"))
+        self.assertTrue(self.trie.starts_with("apple"))
+        self.assertFalse(self.trie.starts_with("apq"))
+
+    def test_multiple_words(self):
+        words = ["apple", "app", "apricot", "banana", "bat"]
+        for word in words:
+            self.trie.insert(word)
+
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+
+        self.assertTrue(self.trie.starts_with("ap"))
+        self.assertTrue(self.trie.starts_with("ba"))
+        self.assertFalse(self.trie.starts_with("c"))
+
+    def test_empty_string(self):
+        self.trie.insert("")
+        self.assertTrue(self.trie.search(""))
+        self.assertTrue(self.trie.starts_with(""))
+
+if __name__ == "__main__":
+    unittest.main()