Solution: Implement Trie (Prefix Tree)

Let's solve the Implement Trie (Prefix Tree) problem using the Trie pattern.

We'll cover the following

Statement

Trie is a tree-like data structure used to store strings. The tries are also called prefix trees because they provide very efficient prefix-matching operations. Implement a trie data structure with three functions that perform the following tasks:

  • Insert (word): This inserts a word into the trie.
  • Search (word): This searches the given word in the trie and returns TRUE, if found. Otherwise, return FALSE.
  • Search prefix (prefix): This searches the given prefix in the trie and returns TRUE, if found. Otherwise, return FALSE.

Constraints:

  • 1≤1 \leq word.length, prefix.length ≤500\leq 500
  • The strings consist only of lowercase English letters.
  • At most, 10310^3 calls in total will be made to the functions.

Pattern: Trie

There are multiple data structures that can be used to store strings. However, for a majority of them, searching a string requires a complete traversal of the stored data. A more efficient approach is to use a trie.

A trie is a tree of characters. The trie takes a word as a parameter and creates a new node for each character of that word. This way, the given string is searched character by character and does not require an exhaustive search. Therefore, this problem fits perfectly in the trie pattern.

Solution

To implement a Trie class, we’ll initialize the root node of the tree in the constructor. This node will be of the Node type. The Node class contains a dictionary of nodes and a boolean variable, which determines if the character is at the end of a word.

Let’s discuss the following functions now:

  • Insert(word): In this function, we take a word as input. We begin from the root node and iterate over the string one character at a time. At each node, we check whether or not a child node with the character is present. If it’s not present, a new node is initialized. For the last character of the word, we also set the boolean variable to TRUE for the corresponding node.

  • Search(word): In this function, we start traversing from the root node and check if the first character is the child of the root node. If so, we move on to that node and check its children for the next character. If, at any point, we do not find the node corresponding to a character, we return FALSE. Alternatively, we return FALSE when we find the whole word, but the boolean variable is not TRUE for the last node, signaling that the word doesn’t end here. We only return TRUE when all the characters match, and the word ends at that point.

  • Search prefix(prefix): This function is mostly the same as the search function. The only exception is that we do not check if the boolean variable is also set to TRUE in the last-found node because we aren’t looking for a complete word but just a prefix. In the trie above, for instance, searchPrefix("by") should return TRUE.

Let’s look at the code for this solution below:

main.java
TrieNode.java
Print.java
import java.util.*;
class Trie {
  TrieNode root;
  // Constructor to create a trie node.
  public Trie() {
      this.root = new TrieNode();
  }
  // A function to insert a word in trie.
  public void insert(String word) {
      TrieNode trieNode = this.root;
      for (int i = 0; i < word.length(); i++) {
          char c = word.charAt(i);
          if (!trieNode.children.containsKey(c)) {
            trieNode.children.put(c, new TrieNode());
          }
          trieNode = trieNode.children.get(c);
      }
      trieNode.isWord = true;
  }
  // A function to search for a word in the trie.
  public boolean search(String word) {
      TrieNode trieNode = this.root;
      
      for (int i = 0; i < word.length(); i++) {
          char c = word.charAt(i);
          if (!trieNode.children.containsKey(c)) {
            return false;
          }
          
        trieNode = trieNode.children.get(c);
      }
      return trieNode.isWord; 
  }
  // A function to search for a prefix of a word in the trie.
  public boolean searchPrefix(String prefix) {
      TrieNode trieNode = this.root;
      
      for (int i = 0; i < prefix.length(); i++) {
          char c = prefix.charAt(i);
          
          if (!trieNode.children.containsKey(c)) {
            return false;
          }
          else 
          trieNode = trieNode.children.get(c); 
      }
      return true;
  }
  
  // Driver Code
  public static void main(String args[]) {
      List < String > keys = Arrays.asList("the", "a", "there", "answer");
      Trie trieOfKeys = new Trie();
      int num = 1;
      for (String x: keys) {
          System.out.println(num + ".\tInserting key: '" + x + "'");
          trieOfKeys.insert(x);
          num += 1;
          System.out.println(new String(new char[100]).replace('\0', '-'));
      }
      List < String > search = Arrays.asList("a", "answer", "xyz", "an");
      for (String y: search) {
          System.out.println(num + ".\tSearching key: '" + y + "'");
          System.out.println("\tKey found? " + trieOfKeys.search(y));
          num += 1;
          System.out.println(new String(new char[100]).replace('\0', '-'));
      }
      List < String > searchPrefix = Arrays.asList("b", "an");
      for (String z: search) {
          System.out.println(num + ".\tSearching prefix: '" + z + "'");
          System.out.println("\tPrefix found? " + trieOfKeys.searchPrefix(z));
          num += 1;
          System.out.println(new String(new char[100]).replace('\0', '-'));
      }
  }
}
Implement Trie (Prefix Tree)

Solution summary

To recap, the solution to this problem can be divided into the following parts:

  • Insert(): Add words to a trie by iterating character by character and checking if each character has a corresponding node. If it does not, a new node is created. The boolean variable is set to TRUE for the last node.

  • Search(): Search by starting from the root and going down the trie, checking character by character. For a complete word, check if the boolean variable is set to TRUE.

  • Search prefix(): Search by starting from the root and going down the trie. However, the boolean variable doesn’t necessarily have to be TRUE.

Time complexity

  • Insert(): The time complexity is O(l)O(l), where ll is the length of the word being inserted.
  • Search(): The time complexity is O(l)O(l), where ll is the length of the word that we need to search in the trie.
  • Search prefix(): The time complexity is O(l)O(l), where ll is the length of the prefix that we need to search in the trie.

Space complexity

  • Insert(): The space complexity is O(l)O(l) because, in the worst case, we will add ll nodes in the trie.
  • Search(): The space complexity is O(1)O(1) because constant space is used while searching the trie.
  • Search prefix(): The space complexity is O(1)O(1) because constant space is used while searching the trie.

Level up your interview prep. Join Educative to access 70+ hands-on prep courses.