Add New Trie Algorithm

tree/trie/new_trie.py

@@ -0,0 +1,97 @@
+"""
+Implement a trie data structure with insert, search, and delete operations.
+
+This implementation provides an efficient tree-like data structure for storing and retrieving strings.
+It is particularly useful for tasks such as prefix matching, autocomplete, and spell checking.
+
+Structure:
+- The trie is composed of TrieNodes, each representing a character in the stored words.
+- Each TrieNode has a dictionary of children nodes and a boolean flag indicating if it's the end of a word.
+
+Functions:
+- insert: Adds a word to the trie by creating a path of nodes for each character.
+- search: Checks if a word exists in the trie by traversing the path of nodes.
+- delete: Removes a word from the trie, cleaning up any unnecessary nodes.
+
+Performance:
+- Time complexity for insert, search, and delete operations is O(m), where m is the length of the word.
+- Space complexity is O(n*m), where n is the number of words and m is the average word length.
+
+Note:
+- This implementation assumes all inputs consist of lowercase letters a-z.
+- For larger datasets or more complex use cases, additional optimizations may be necessary.
+"""
+
+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.
+        """
+        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.
+        """
+        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
+
+    def delete(self, word: str) -> None:
+        """
+        Removes a word from the trie if it exists.
+        """
+        def _delete_helper(node, word, index):
+            if index == len(word):
+                if not node.is_end_of_word:
+                    return False
+                node.is_end_of_word = False
+                return len(node.children) == 0
+
+            char = word[index]
+            if char not in node.children:
+                return False
+
+            should_delete_current_node = _delete_helper(node.children[char], word, index + 1)
+
+            if should_delete_current_node:
+                del node.children[char]
+                return len(node.children) == 0
+
+            return False
+
+        _delete_helper(self.root, word, 0)
+
+# Example usage
+if __name__ == "__main__":
+    trie = Trie()
+    words = ["apple", "app", "apricot", "banana"]
+    for word in words:
+        trie.insert(word)
+
+    print(trie.search("apple"))  # True
+    print(trie.search("app"))    # True
+    print(trie.search("apricot"))  # True
+    print(trie.search("banana"))   # True
+    print(trie.search("grape"))    # False
+
+    trie.delete("apple")
+    print(trie.search("apple"))  # False
+    print(trie.search("app"))    # True
+    print(trie.search("apricot"))  # True

tree/trie/test_new_trie.py

@@ -0,0 +1,60 @@
+import unittest
+from new_trie import NewTrie
+
+class TestNewTrie(unittest.TestCase):
+    def setUp(self):
+        self.trie = NewTrie()
+
+    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_empty_string(self):
+        self.trie.insert("")
+        self.assertTrue(self.trie.search(""))
+
+    def test_common_prefix(self):
+        self.trie.insert("cat")
+        self.trie.insert("car")
+        self.assertTrue(self.trie.search("cat"))
+        self.assertTrue(self.trie.search("car"))
+        self.assertFalse(self.trie.search("ca"))
+
+    def test_non_existent_word(self):
+        self.trie.insert("hello")
+        self.assertFalse(self.trie.search("world"))
+
+    def test_words_with_common_prefixes(self):
+        words = ["tree", "trie", "try", "triangle", "trap"]
+        for word in words:
+            self.trie.insert(word)
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+        self.assertFalse(self.trie.search("tr"))
+        self.assertFalse(self.trie.search("tria"))
+
+    def test_prefix_search(self):
+        self.trie.insert("programming")
+        self.assertTrue(self.trie.search("programming"))
+        self.assertFalse(self.trie.search("program"))
+        self.assertFalse(self.trie.search("prog"))
+
+    def test_single_character_words(self):
+        for char in "abcdefghijklmnopqrstuvwxyz":
+            self.trie.insert(char)
+        for char in "abcdefghijklmnopqrstuvwxyz":
+            self.assertTrue(self.trie.search(char))
+        self.assertFalse(self.trie.search(""))
+
+    def test_similar_non_existent_words(self):
+        self.trie.insert("hello")
+        self.trie.insert("world")
+        self.assertFalse(self.trie.search("hell"))
+        self.assertFalse(self.trie.search("worl"))
+        self.assertFalse(self.trie.search("he"))
+        self.assertFalse(self.trie.search("wo"))
+
+if __name__ == "__main__":
+    unittest.main()