Skip to content

Commit 4fa6200

Browse files
committed
added trie with additional features
1 parent c3d3b74 commit 4fa6200

File tree

2 files changed

+233
-0
lines changed

2 files changed

+233
-0
lines changed
Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
import pytest
2+
from trie import Trie
3+
4+
def test_trie_insert_search():
5+
trie = Trie()
6+
trie.insert("apple")
7+
assert trie.search("apple")
8+
assert not trie.search("app")
9+
trie.insert("app")
10+
assert trie.search("app")
11+
12+
def test_trie_starts_with():
13+
trie = Trie()
14+
trie.insert("apple")
15+
assert trie.starts_with("app")
16+
assert trie.starts_with("a")
17+
assert not trie.starts_with("b")
18+
assert not trie.starts_with("applxyz")
19+
20+
def test_trie_empty():
21+
trie = Trie()
22+
assert not trie.search("apple")
23+
assert not trie.starts_with("app")
24+
25+
def test_trie_multiple_words():
26+
trie = Trie()
27+
trie.insert("apple")
28+
trie.insert("application")
29+
trie.insert("banana")
30+
assert trie.search("apple")
31+
assert trie.search("application")
32+
assert trie.search("banana")
33+
assert not trie.search("app")
34+
assert trie.starts_with("app")
35+
assert trie.starts_with("ban")
36+
assert not trie.starts_with("aplx")
37+
38+
def test_trie_case_sensitive():
39+
trie = Trie()
40+
trie.insert("Apple")
41+
assert trie.search("Apple")
42+
assert not trie.search("apple")
43+
44+
def test_count_words():
45+
trie = Trie()
46+
assert trie.count_words() == 0
47+
trie.insert("apple")
48+
assert trie.count_words() == 1
49+
trie.insert("app")
50+
assert trie.count_words() == 2
51+
trie.insert("apple")
52+
assert trie.count_words() == 2
53+
54+
def test_longest_common_prefix():
55+
trie = Trie()
56+
assert trie.longest_common_prefix() == ""
57+
trie.insert("apple")
58+
assert trie.longest_common_prefix() == "apple"
59+
trie.insert("application")
60+
assert trie.longest_common_prefix() == "appl"
61+
trie.insert("banana")
62+
assert trie.longest_common_prefix() == ""
63+
64+
def test_autocomplete():
65+
trie = Trie()
66+
trie.insert("apple")
67+
trie.insert("application")
68+
trie.insert("app")
69+
assert trie.autocomplete("app") == ["app", "apple", "application"]
70+
assert trie.autocomplete("appl") == ["apple", "application"]
71+
assert trie.autocomplete("b") == []
72+
73+
def test_bulk_insert():
74+
trie = Trie()
75+
trie.bulk_insert(["apple", "banana", "orange"])
76+
assert trie.search("apple")
77+
assert trie.search("banana")
78+
assert trie.search("orange")
79+
assert trie.count_words() == 3
80+
81+
def test_clear():
82+
trie = Trie()
83+
trie.insert("apple")
84+
trie.clear()
85+
assert trie.is_empty()
86+
assert trie.count_words() == 0
87+
assert not trie.search("apple")
88+
89+
def test_is_empty():
90+
trie = Trie()
91+
assert trie.is_empty()
92+
trie.insert("apple")
93+
assert not trie.is_empty()
94+
trie.clear()
95+
assert trie.is_empty()
96+
97+
def test_find_all_words():
98+
trie = Trie()
99+
trie.bulk_insert(["apple", "banana", "orange"])
100+
assert sorted(trie.find_all_words()) == sorted(["apple", "banana", "orange"])
101+
trie.clear()
102+
assert trie.find_all_words() == []
103+
104+
105+
def test_longest_word():
106+
trie = Trie()
107+
assert trie.longest_word() is None
108+
trie.bulk_insert(["apple", "banana", "application"])
109+
assert trie.longest_word() == "application"
110+
trie.insert("a")
111+
assert trie.longest_word() == "application"

pydatastructs/trees/trie.py

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# trie.py
2+
3+
class TrieNode:
4+
"""Represents a node in the Trie data structure."""
5+
def __init__(self):
6+
"""Initializes a TrieNode with empty children and is_end_of_word set to False."""
7+
self.children = {}
8+
self.is_end_of_word = False
9+
self.word = None
10+
11+
class Trie:
12+
"""Represents the Trie (prefix tree) data structure."""
13+
def __init__(self):
14+
"""Initializes an empty Trie with a root TrieNode."""
15+
self.root = TrieNode()
16+
self.word_count = 0
17+
18+
def insert(self, word):
19+
"""Inserts a word into the Trie."""
20+
node = self.root
21+
for char in word:
22+
if char not in node.children:
23+
node.children[char] = TrieNode()
24+
node = node.children[char]
25+
if not node.is_end_of_word:
26+
node.is_end_of_word = True
27+
node.word = word
28+
self.word_count += 1
29+
30+
def search(self, word):
31+
"""Searches for a word in the Trie."""
32+
node = self.root
33+
for char in word:
34+
if char not in node.children:
35+
return False
36+
node = node.children[char]
37+
return node.is_end_of_word
38+
39+
def starts_with(self, prefix):
40+
"""Checks if any word in the Trie starts with the given prefix."""
41+
node = self.root
42+
for char in prefix:
43+
if char not in node.children:
44+
return False
45+
node = node.children[char]
46+
return True
47+
48+
def count_words(self):
49+
"""Returns the total number of words stored in the Trie."""
50+
return self.word_count
51+
52+
def longest_common_prefix(self):
53+
"""Finds the longest common prefix among all words in the Trie."""
54+
node = self.root
55+
prefix = ""
56+
while len(node.children) == 1 and not node.is_end_of_word:
57+
char = next(iter(node.children))
58+
prefix += char
59+
node = node.children[char]
60+
return prefix
61+
62+
def autocomplete(self, prefix):
63+
"""Provides a list of words that match a given prefix."""
64+
node = self.root
65+
for char in prefix:
66+
if char not in node.children:
67+
return []
68+
node = node.children[char]
69+
70+
def collect_words(current_node, current_prefix):
71+
words = []
72+
if current_node.is_end_of_word:
73+
words.append(current_prefix)
74+
for char, child_node in current_node.children.items():
75+
words.extend(collect_words(child_node, current_prefix + char))
76+
return words
77+
78+
return collect_words(node, prefix)
79+
80+
def bulk_insert(self, words):
81+
"""Inserts multiple words into the Trie in a single operation."""
82+
for word in words:
83+
self.insert(word)
84+
85+
def clear(self):
86+
"""Removes all words from the Trie, resetting it."""
87+
self.root = TrieNode()
88+
self.word_count = 0
89+
90+
def is_empty(self):
91+
"""Returns True if the Trie is empty, otherwise False."""
92+
return self.word_count == 0
93+
94+
def find_all_words(self):
95+
"""Retrieves all words currently stored in the Trie."""
96+
def collect_words(current_node):
97+
words = []
98+
if current_node.is_end_of_word:
99+
words.append(current_node.word)
100+
for child_node in current_node.children.values():
101+
words.extend(collect_words(child_node))
102+
return words
103+
104+
return collect_words(self.root)
105+
106+
def shortest_unique_prefix(self, word):
107+
"""Determines the shortest unique prefix for a given word."""
108+
node = self.root
109+
prefix = ""
110+
for char in word:
111+
prefix += char
112+
if len(node.children[char].children) <= 1 and node.children[char].is_end_of_word == False :
113+
return prefix
114+
node = node.children[char]
115+
return word
116+
117+
def longest_word(self):
118+
"""Finds and returns the longest word in the Trie."""
119+
all_words = self.find_all_words()
120+
if not all_words:
121+
return None
122+
return max(all_words, key=len)

0 commit comments

Comments
 (0)