diff --git a/trie/test_trie.py b/trie/test_trie.py new file mode 100644 index 0000000..03856fa --- /dev/null +++ b/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() diff --git a/trie/trie.py b/trie/trie.py new file mode 100644 index 0000000..05912f7 --- /dev/null +++ b/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