from trie_node import *
class WordDictionary:
# Initialize the root with TrieNode and set
# the 'can_find' boolean to FALSE
def __init__(self):
self.root = TrieNode()
self.can_find = False
# Function to add a new word to the dictionary
def add_word(self, word):
n = len(word)
cur_node = self.root
for i, val in enumerate(word):
index = ord(val) - ord('a')
if cur_node.children[index] is None:
cur_node.children[index] = TrieNode()
cur_node = cur_node.children[index]
if i == n - 1:
if cur_node.complete:
print("\tWord already present!")
return
cur_node.complete = True
print("\tWord added successfully!")
# Function to search for a word in the dictionary
def search_word(self, word):
self.can_find = False
self.search_helper(self.root, word, 0)
return self.can_find
def search_helper(self, node, word, i):
if self.can_find:
return
if not node:
return
if len(word) == i:
if node.complete:
self.can_find = True
return
if word[i] == '.':
for j in range(ord('a'), ord('z') + 1):
self.search_helper(node.children[j - ord('a')], word, i + 1)
else:
index = ord(word[i]) - ord('a')
self.search_helper(node.children[index], word, i + 1)
# Function to get all words in the dictionary
def get_words(self):
words_list = []
if not self.root:
return []
return self.dfs(self.root, "", words_list)
def dfs(self, node, word, words_list):
if not node:
return words_list
if node.complete:
words_list.append(word)
for j in range(ord('a'), ord('z') + 1):
prefix = word + chr(j)
words_list = self.dfs(node.children[j - ord('a')], prefix, words_list)
return words_list