Implement Trie Data Structure

trie/test_trie.py

@@ -0,0 +1,188 @@
+import unittest
+from trie import Trie
+
+class TestTrie(unittest.TestCase):
+    def setUp(self):
+        self.trie = Trie()
+
+    def test_insert_and_search(self):
+        words = ["apple", "app", "apricot", "banana"]
+        for word in words:
+            self.trie.insert(word)
+
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+
+        self.assertFalse(self.trie.search("grape"))
+        self.assertFalse(self.trie.search("ap"))
+
+    def test_delete(self):
+        words = ["apple", "app", "apricot", "banana"]
+        for word in words:
+            self.trie.insert(word)
+
+        self.assertTrue(self.trie.delete("apple"))
+        self.assertFalse(self.trie.search("apple"))
+        self.assertTrue(self.trie.search("app"))
+
+        self.assertTrue(self.trie.delete("banana"))
+        self.assertFalse(self.trie.search("banana"))
+
+        self.assertFalse(self.trie.delete("grape"))
+
+    def test_empty_string(self):
+        self.trie.insert("")
+        self.assertTrue(self.trie.search(""))
+        self.assertTrue(self.trie.delete(""))
+        self.assertFalse(self.trie.search(""))
+
+    def test_long_word(self):
+        long_word = "a" * 1000
+        self.trie.insert(long_word)
+        self.assertTrue(self.trie.search(long_word))
+        self.assertTrue(self.trie.delete(long_word))
+        self.assertFalse(self.trie.search(long_word))
+
+    def test_common_prefixes(self):
+        words = ["car", "card", "carpet", "carpool", "cat"]
+        for word in words:
+            self.trie.insert(word)
+
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+
+        self.assertTrue(self.trie.search("car"))
+        self.assertFalse(self.trie.search("carp"))
+
+    def test_delete_common_prefixes(self):
+        words = ["car", "card", "carpet"]
+        for word in words:
+            self.trie.insert(word)
+
+        self.assertTrue(self.trie.delete("card"))
+        self.assertFalse(self.trie.search("card"))
+        self.assertTrue(self.trie.search("car"))
+        self.assertTrue(self.trie.search("carpet"))
+
+    def test_prefix_search(self):
+        words = ["apple", "app", "application", "apply"]
+        for word in words:
+            self.trie.insert(word)
+
+        self.assertTrue(self.trie.search("app"))
+        self.assertTrue(self.trie.search("apple"))
+        self.assertFalse(self.trie.search("ap"))
+
+    def test_single_character_words(self):
+        words = ["a", "b", "c", "ab", "bc"]
+        for word in words:
+            self.trie.insert(word)
+
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+
+        self.assertTrue(self.trie.delete("a"))
+        self.assertFalse(self.trie.search("a"))
+        self.assertTrue(self.trie.search("ab"))
+
+    def test_case_sensitivity(self):
+        words = ["Apple", "apple", "APPLE", "bAnAnA"]
+        for word in words:
+            self.trie.insert(word)
+
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+
+        self.assertFalse(self.trie.search("aPpLe"))
+        self.assertTrue(self.trie.delete("Apple"))
+        self.assertFalse(self.trie.search("Apple"))
+        self.assertTrue(self.trie.search("apple"))
+
+    def test_special_characters(self):
+        words = ["hello!", "@world", "#python", "$100"]
+        for word in words:
+            self.trie.insert(word)
+
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+
+        self.assertFalse(self.trie.search("hello"))
+        self.assertFalse(self.trie.search("world"))
+
+    def test_complex_operations(self):
+        words = ["tree", "trie", "algo", "algorithm", "algorithms"]
+        for word in words:
+            self.trie.insert(word)
+
+        self.assertTrue(self.trie.search("algo"))
+        self.assertTrue(self.trie.delete("algorithm"))
+        self.assertFalse(self.trie.search("algorithm"))
+        self.assertTrue(self.trie.search("algorithms"))
+
+        self.trie.insert("algorithmic")
+        self.assertTrue(self.trie.search("algorithmic"))
+        self.assertFalse(self.trie.search("algorithm"))
+
+    def test_unicode_characters(self):
+        words = ["café", "résumé", "über", "naïve"]
+        for word in words:
+            self.trie.insert(word)
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+        self.assertFalse(self.trie.search("cafe"))
+        self.assertFalse(self.trie.search("resume"))
+
+    def test_mixed_operations(self):
+        self.trie.insert("programming")
+        self.trie.insert("program")
+        self.assertTrue(self.trie.search("program"))
+        self.trie.delete("programming")
+        self.assertFalse(self.trie.search("programming"))
+        self.assertTrue(self.trie.search("program"))
+        self.trie.insert("progress")
+        self.assertTrue(self.trie.search("progress"))
+
+    def test_empty_trie(self):
+        self.assertFalse(self.trie.search("any"))
+        self.assertFalse(self.trie.delete("any"))
+
+    def test_large_dataset(self):
+        words = [f"word{i}" for i in range(10000)]
+        for word in words:
+            self.trie.insert(word)
+        for word in words:
+            self.assertTrue(self.trie.search(word))
+        self.assertFalse(self.trie.search("nonexistent"))
+
+    def test_prefix_matching(self):
+        words = ["prefix", "preface", "prepare", "prevent"]
+        for word in words:
+            self.trie.insert(word)
+
+        # 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"))
+        self.trie.delete("a")
+        self.assertFalse(self.trie.search("a"))
+        self.trie.insert("ab")
+        self.assertTrue(self.trie.search("ab"))
+        self.assertFalse(self.trie.search("a"))
+
+if __name__ == "__main__":
+    unittest.main()

trie/trie.py

@@ -0,0 +1,100 @@
+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):
+        """
+        Insert a word into the trie.
+        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:
+                node.children[char] = TrieNode()
+            node = node.children[char]
+        node.is_end_of_word = True
+
+    def search(self, word, is_prefix=False):
+        """
+        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 is_prefix or node.is_end_of_word
+
+    def delete(self, word):
+        """
+        Delete a word from the trie.
+        Returns True if the word was deleted, False if it wasn't found.
+        Time complexity: O(m), where m is the length of the word.
+        """
+        if not word:
+            if self.root.is_end_of_word:
+                self.root.is_end_of_word = False
+                return True
+            return False
+
+        stack = [(self.root, 0)]
+        last_node_with_branch = None
+        last_index_with_branch = 0
+
+        while stack:
+            node, index = stack.pop()
+
+            if index == len(word):
+                if not node.is_end_of_word:
+                    return False
+                node.is_end_of_word = False
+                break
+
+            char = word[index]
+            if char not in node.children:
+                return False
+
+            if len(node.children) > 1 or node.is_end_of_word:
+                last_node_with_branch = node
+                last_index_with_branch = index
+
+            stack.append((node.children[char], index + 1))
+
+        if not node.children:
+            if last_node_with_branch:
+                del last_node_with_branch.children[word[last_index_with_branch]]
+            else:
+                self.root.children.pop(word[0], None)
+
+        return True
+
+# Example usage
+if __name__ == "__main__":
+    trie = Trie()
+
+    # Insert words
+    words = ["apple", "app", "apricot", "banana"]
+    for word in words:
+        trie.insert(word)
+
+    # Search words
+    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
+
+    # Delete words
+    print(trie.delete("apple"))    # True
+    print(trie.search("apple"))    # False
+    print(trie.search("app"))      # True
+
+    print(trie.delete("banana"))   # True
+    print(trie.search("banana"))   # False