Trie implementation [Python & Java]

Trie_Node_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 Trie_Python_Java


Corresponding Trie structure Trie_Tree_Structure_Python_Java


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


Copyright (c) 2019-2021, Algotree.org.
All rights reserved.