How to implement Trie Data Structure?

Trie is a datastructure useful to store set of words.

It is a tree with a lot of children , one for each letter of the alphabets.

So if you want to store lower case English letter words each node can have unto 26 child nodes.

Every word is represented by as many nodes as there are number of letters in the word with each letter node connected to each other.

For example , the word cat be represented as:

There will always be a singe root node to which the words are connected.

Now if you want to add the word cap to the above Trie , you would do it this way:

As you see , since the first two letters of the two words cat and cap are the same they are represented by the same nodes c and a connected to the root.

Then we add p to the node a to represent the word cap.

So to add a word to a Trie,

We see if the first letter of the word is already connected to the root ,

If yes we will keep moving down until the letters match and then append the remaining letters as separate connected nodes.

If it is not present then we will append a new letter chain to the root node.

This structure lets us easily search for a word and add new words.

It is faster and occupies less space.

This data structure is used in use cases like auto complete feature which you see in Google search , T9 dictionary , spell checker etc.

Let’s see how to implement this Trie data structure.

Try out the solution here:

https://leetcode.com/problems/design-add-and-search-words-data-structure/

Data Structure of Trie

First let’s create the Trie data structure.

We need :

  • a set of children as many as 26 for each letter in the lower case English alphabets
  • a boolean variable to indicate if we have reached the end of a word – this will help in searching for a word
  • a put operation to add a letter node to the trie to the current node
  • a get operation to get a Trie node for the given character
  • a setEnd() method to mark the end of a word
  • a containsKey() method to check if a given character is present in the children of the current node.

Here is the code:

class Trie{
    
    
    Trie[] links = new Trie[26];
  
    
    
    boolean isEnd;
    
    public void put(char c,Trie node){
        
        
        links[c - 'a'] = node;
        
        
    }
    
    public Trie get(char c){
        
        return links[c - 'a'];
    }
    
    
    public boolean containsKey(char c){
        
        
        return links[c - 'a'] !=null;
    }
    
    public void setEnd(){
        
        isEnd = true;
    }
    
}

Now let’s implement the operation to add a new word to the Trie data structure.

For this you iterate through each character in the word ,

And see if the current node in the Trie contains the character , if so you move to that node or else you create a node for the current character and add to the current node and move to that node.

Once you reach the end , you mark the last node as the end node.

Adding a word to Trie

Here is the method implementation:

    public void addWord(String word) {
        
        
        Trie node = root;
        
        for(int i=0;i<word.length();i++){
            
            char c = word.charAt(i);
            if(!node.containsKey(c)){
                
                node.put(c,new Trie());
            }
            
            node = node.get(c);
        }
        
        node.setEnd();
        
    }
    

Searching for a word:

For searching again you check if the first character is a child to the root node,

If so you move to that node.

Then you check if the second character is a child of the previous node .

Then you move to that node and so on until you reach the last character .

If the last character is marked as the end node then we have found our word.

At any stage in the previous steps if we don’t find the next character we return false since the character chain is not present in the Trie.

    
    public boolean search(String word) {
        
        TrieNode node = root;
        for(int i=0;i<word.length();i++){
            
            
            char c = word.charAt(i);
            
            if(node.containsKey(c)){
                
                node = node.get(c);
            }else{
                
                return false;
            }
            
        }
        
        return node!=null && node.isEnd;
        
        
    }

Searching for a prefix:

Searching for a prefix is the same as searching for a word except that you don’t check if the last character in the prefix is the end node:

    public boolean startsWith(String prefix) {
        
TrieNode node = root;
        
        for(int i=0;i<prefix.length();i++){
            
            
            char c = prefix.charAt(i);
            
            if(node.containsKey(c)){
                
                node = node.get(c);
            }else{
                
                return false;
            }
        }
        
        return node!=null;
        
    }

Here is the entire code:

class Trie {

    
    TrieNode root ;
    public Trie() {
        
    root = new TrieNode();
        
    }
    
    public void insert(String word) {
        
        TrieNode node = root;
        for(int i=0;i<word.length();i++){
            
            char c = word.charAt(i);
            
            
            if(!node.containsKey(c)){
                
                node.put(c,new TrieNode());
            }
            
            node = node.get(c);
        }
        
        node.setEnd();
        
        
    }
    
    public boolean search(String word) {
        
        TrieNode node = root;
        for(int i=0;i<word.length();i++){
            
            
            char c = word.charAt(i);
            
            if(node.containsKey(c)){
                
                node = node.get(c);
            }else{
                
                return false;
            }
            
        }
        
        return node!=null && node.isEnd;
        
        
    }
    
    public boolean startsWith(String prefix) {
        
TrieNode node = root;
        
        for(int i=0;i<prefix.length();i++){
            
            
            char c = prefix.charAt(i);
            
            if(node.containsKey(c)){
                
                node = node.get(c);
            }else{
                
                return false;
            }
        }
        
        return node!=null;
        
    }
}

class TrieNode{
    
    
    TrieNode[] links = new TrieNode[26];
    
    boolean isEnd;
    
    public boolean containsKey(char key){
        
        return links[key - 'a'] != null;
    
    }
    
    public TrieNode get(char c){
        
        return links[c -'a'];
    }
    
    public void put(char c , TrieNode root){
        
        
        links[c -'a'] = root;
    }
    
    public void setEnd(){
        
        isEnd = true;
    }
    
}
/**
 * Your Trie object will be instantiated and called as such:
 * Trie obj = new Trie();
 * obj.insert(word);
 * boolean param_2 = obj.search(word);
 * boolean param_3 = obj.startsWith(prefix);
 */

Time complexity:

Adding a word can take O(m) time complexity where m is the number of letters in the word .

Searching a word also takes the same time complexity since you need to navigate through each letter node.

Space complexity:

Storing a word requires m nodes where m is the number of letters in the word. So space complexity is O(m).

Searching a word doesn’t require any extra space complexity ,so space complexity is O(1)

That’s it!

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s