Trie implementation [Python & Java]
Trie or Prefix Tree / Radix Tree
Trie and Trie Node
-
The data structure of a trie resembles a tree.
-
Each node in the trie stores the below:
- A dictionary (Python) / HashMap (Java) of character as key and a corresponding Trie node as a value.
- A character for storing a letter.
- A boolean flag that indicates if the word ends with the stored letter. In the below example the green nodes mark the end of the word.
-
In the example, the trie has the capacity for storing
all the english words containing small letters . -
Words are retrieved by traversing from the root of the trie down the branch of the tree.
Trie for storing lower case english words
Corresponding Trie structure
Inserting a string into a Trie
a) Start from the root node of the Trie i.e current node = root; do the below
b) For every letter ‘c’ in the string
c) If there is a key (letter ‘c’) in the dictionary / map,
i.e the substring till the letter c is already inserted into the trie which also means
that the current node is the parent of the child node pointed at by the letter ‘c’.
Point the current node to the child node for inserting the next letter.
d) Else
Create an entry of letter ‘c’ and corresponding child node to be inserted into the TRIE.
Point current node to the child node in the dictionary / map for inserting the next letter.
e) Set the end of the word for the last TRIE node as the last letter in the string has been inserted.
Searching a string in a Trie
a) Start from the root node of the Trie i.e current_node = root; do the below
b) For every letter ‘c’ in the string
c) If there is key (letter ‘c’) in the dictionary / map,
i.e the substring till the letter c is present in the Trie which also means
that the current node is the parent of the child node pointed at by the letter ‘c’.
Point current node to the child node for searching the next letter.
d) Else
None of the children of the current node has the letter ‘c’ that we are searching for.
Set the current node to None and stop the search.
e) Return current_node.
Printing words with a given prefix in a lexical order in a Trie.
a) Search for the node that indicates the end of the prefix. Use the search routine described above for achieving this.
The node found is the current node.
b) Use a recursive function with current_node, given prefix and suffix set to an empty string.
PrintLexical (current_node, prefix, suffix)
c) If the current_node has the end_of_the_word set and suffix size if greater than 0 then
d) Print the string prefix + suffix.
e) For each child of the current node do the following.
f) Set string temp = suffix as there could be more words extending the earlier found suffix in the tree.
g) Append the letter of the child to the temp string.
i.e temp = temp + letter_of_the_child.
h) Call PrintLexical (current_node, prefix, temp).
Time complexity of creating a TRIE : O(N*L) where ‘N’ words are to be inserted into the trie and the average length of each word is ‘L’. Time complexity of searching in a TRIE : O(N*L) where ‘N’ words are to be searched in the trie and the average length of each word is ‘L’.
Implementation of Trie insert, search and print operations
class TrieNode:
"""A trie node"""
def __init__(self, char):
# Character stored in this node
self.char = char
# A flag that marks if the word ends on this particular node.
self.end_of_word = False
# A dictionary of child nodes where the keys are the characters (letters)
# and values are the nodes
self.children = {}
class Trie:
"""The trie structure"""
def __init__(self):
"""
The root node of a trie is empty and does not store any character.
"""
self.root = TrieNode("")
def Insert (self, string):
"""Insert a string i.e a word into the trie"""
node = self.root
# Check each character in the string
# If none of the children of the current node contains the character,
# create a new child of the current node for storing the character.
for char in string:
if char in node.children:
node = node.children[char]
else:
# As the character is not found, create a new trie node
new_node = TrieNode(char)
node.children[char] = new_node
node = new_node
# Mark the end of a word
node.end_of_word = True
def Search (self, string):
"""Search a string i.e search a word in the trie"""
node = self.root
# Check each character in the string
# If none of the children of the node contains the character,
# Return none
for char in string:
if char in node.children:
node = node.children[char]
else:
node = None
break
return node
def PrintLexical (self, node, prefix, suffix):
"""Print words in the Trie in lexical order"""
if node.end_of_word == True and len(suffix):
print (prefix + suffix)
# Iterate through all the nodes in the dictionary
# key is the character and children[key] is the child node
for key in node.children:
temp = suffix
child = node.children[key]
temp += key
self.PrintLexical (child, prefix, temp)
def main():
t = Trie()
t.Insert("we")
t.Insert("walk")
t.Insert("want")
t.Insert("wish")
t.Insert("wit")
t.Insert("am")
t.Insert("yo")
t.Insert("will")
t.Insert("wee")
t.Insert("war")
t.Insert("win")
t.Insert("warp")
# Search for a prefix in the Trie
prefix = input("Enter prefix : ")
node = t.Search(prefix)
if (node == None or node == t.root):
print("No words with matching prefix found")
else:
# Prefix has been found in the Trie. Now look for children
if len(node.children):
print("Words with prefix : " + prefix)
t.PrintLexical(node, prefix, "")
else:
print("Just the prefix exists in the Trie, but not the matching words")
if __name__ == "__main__" :
main()
Output
Enter prefix : w
Words with prefix : w
we
wee
walk
want
war
warp
wish
wit
will
win
Enter prefix : wa
Words with prefix : wa
walk
want
war
warp
Enter prefix : war
Words with prefix : war
warp
Enter prefix : warp
Just the prefix exist in the Trie, but not the matching words
import java.util.Scanner;
import java.util.HashMap;
class TrieNode {
public char letter;
public boolean end_of_word; // A flag that marks if the word ends on this particular node.
public HashMap <Character, TrieNode> children;
TrieNode (char c) {
// Character stored in this node
letter = c;
end_of_word = false;
// A dictionary of child nodes where the keys are the characters (letters)
// and values are the nodes
children = new HashMap<Character, TrieNode>();
}
}
class Trie {
TrieNode root;
Trie () {
// The root node of a trie is empty and does not store any character.
root = new TrieNode('0');
}
// Insert a string i.e a word into the trie
public void Insert (String word) {
TrieNode node = root;
// Check each character in the string
// If none of the children of the current node contains the character,
// create a new child of the current node for storing the character.
for (int i = 0; i < word.length(); i++) {
char c = word.charAt(i);
if (node.children.containsKey(c)) {
node = node.children.get(c);
} else {
// As the character is not found, create a new trie node
TrieNode new_node = new TrieNode(c);
node.children.put(c, new_node);
node = new_node;
}
}
// Mark the end of a word
node.end_of_word = true;
}
// Search a string i.e search a word in the trie"""
public TrieNode Search (String word) {
TrieNode node = root;
// Check each character in the string
// If none of the children of the node contains the character,
// Return none
for (int i = 0; i < word.length(); i++) {
char c = word.charAt(i);
if (node.children.containsKey(c)) {
node = node.children.get(c);
} else {
node = null;
break;
}
}
return node;
}
// Print words in the Trie in lexical order"""
public void PrintLexical (TrieNode node, String prefix, String suffix) {
if (node.end_of_word == true && suffix.length() > 0) {
System.out.println (prefix + suffix);
}
// Iterate through all the nodes in the dictionary
// key is the character and children[key] is the child node
for (char key : node.children.keySet()) {
String temp = suffix;
TrieNode child = node.children.get(key);
temp += key;
PrintLexical (child, prefix, temp);
}
}
public static void main (String args[]) {
Trie t = new Trie();
t.Insert("we");
t.Insert("walk");
t.Insert("want");
t.Insert("wish");
t.Insert("wit");
t.Insert("am");
t.Insert("yo");
t.Insert("will");
t.Insert("wee");
t.Insert("war");
t.Insert("win");
t.Insert("warp");
// Search for a prefix in the Trie
Scanner scanner_obj = new Scanner(System.in); // Create a Scanner object
System.out.print("Enter prefix : ");
String prefix = scanner_obj.nextLine();
TrieNode node = t.Search(prefix);
if (node == null || node == t.root) {
System.out.println("No words with matching prefix found");
} else {
// Prefix has been found in the Trie. Now look for children
if (node.children.size() > 0) {
System.out.println("Words with prefix : " + prefix);
t.PrintLexical(node, prefix, "");
} else {
System.out.println("Just the prefix exists in the Trie, but not the matching words");
}
}
}
}
Output
Enter prefix : w
Words with prefix : w
we
wee
walk
want
war
warp
wish
wit
will
win
Enter prefix : wa
Words with prefix : wa
walk
want
war
warp
Enter prefix : war
Words with prefix : war
warp
Enter prefix : warp
Just the prefix exist in the Trie, but not the matching words
Enter prefix : x
No words with matching prefix found