Count the number of words with given prefix using Trie

Prerequisite: Trie

Given a list of string str[] and a prefix string pre. The task is to count the number of words in the list of string with a given prefix using trie.

Examples:

Input: str = [ “apk”, “app”, “apple”, “arp”, “array” ], pre = “ap” 
Output: 3 
Explanation:
Below is the representation of trie from using above string. 

The words in str having prefix “ap” are apk, app and apple. 
So, the count is 3

Input: str = [ “gee”, “geek”, “geezer”, “geeksforgeeks”, “geekiness”, “geekgod” ], pre = “geek” 
Output: 4 

Approach:
To solve this problem Trie Data Structure is used and each node of this Trie contains the following three fields:  

1. children: This field is used for mapping from a character to the next level trie node
2. isEndOfWord: This field is used to distinguish the node as the end of the word node
3. num: This field is used to count the number of times a node is visited during insertion in trie

Steps:

• Insert a list of string in trie such that every string in the list is inserted as an individual trie node.
• During inserting update all the fields in every node of the trie
• For a given prefix, traverse the trie till we reach the last character of the given prefix pre.
• The value of the num field in the last node of string pre is the count of prefix in the given list of string.

Below is the implementation of the above approach:  

C++14

// C++ implementation of counting the
// number of words in a trie with a
// given prefix
#include "bits/stdc++.h"
using namespace std;

// Trie Node
struct TrieNode {

    // Using map to store the pointers
    // of children nodes for dynamic
    // implementation, for making the
    // program space efficient
    map<char, TrieNode*> children;

    // If isEndOfWord is true, then
    // node represents end of word
    bool isEndOfWord;

    // num represents number of times
    // a character has appeared during
    // insertion of the words in the
    // trie
    map<char, int> num;
};

// Declare root node
struct TrieNode* root;

// Function to create New Trie Node
struct TrieNode* getNewTrieNode()
{
    struct TrieNode* pNode = new TrieNode;
    pNode->isEndOfWord = false;
    return pNode;
}

// Function to insert a string in trie
void insertWord(string word)
{
    // To hold the value of root
    struct TrieNode* current = root;

    // To hold letters of the word
    char s;

    // Traverse through strings in list
    for (int i = 0; i < word.length(); i++) {
        s = word[i];

        // If s is not present in the
        // character field of current node
        if (current->children.find(s)
            == current->children.end()) {

            // Get new node
            struct TrieNode* p = getNewTrieNode();

            // Insert s in character
            // field of current node
            // with reference to node p
            (current->children)[s] = p;

            // Insert s in num field
            // of current node with
            // value 1
            (current->num)[s] = 1;
        }
        else {

            // Increment the count
            // corresponding to the
            // character s
            current->num[s] = (current->num)[s] + 1;
        }

        // Go to next node
        current = (current->children)[s];
    }
    current->isEndOfWord = true;
}

// Function to count the number of
// words in trie with given prefix
int countWords(vector<string>& words,
               string prefix)
{
    root = getNewTrieNode();

    // Size of list of string
    int n = words.size();

    // Construct trie containing
    // all the words
    for (int i = 0; i < n; i++) {
        insertWord(words[i]);
    }

    struct TrieNode* current = root;
    char s;

    // Initialize the wordCount = 0
    int wordCount = 0;

    for (int i = 0; prefix[i]; i++) {
        s = prefix[i];

        // If the complete prefix is
        // not present in the trie
        if (current->children.find(s)
            == current->children.end()) {

            // Make wordCount 0 and
            // break out of loop
            wordCount = 0;
            break;
        }

        // Update the wordCount
        wordCount = (current->num)[s];

        // Go to next node
        current = (current->children)[s];
    }

    return wordCount;
}

// Driver Code
int main()
{
    // input list of words
    vector<string> words;
    words = { "apk", "app", "apple",
              "arp", "array" };

    // Given prefix to find
    string prefix = "ap";

    // Print the number of words with
    // given prefix
    cout << countWords(words, prefix);
    return 0;
}

Java

// Java implementation of counting the
// number of words in a trie with a
// given prefix
import java.util.HashMap;
import java.util.Map;
import java.util.Vector;

// Trie Node
class TrieNode
{

  // Using map to store the pointers
  // of children nodes for dynamic
  // implementation, for making the
  // program space efficient
  Map<Character, TrieNode> children;

  // If isEndOfWord is true, then
  // node represents end of word
  boolean isEndOfWord;

  // num represents number of times
  // a character has appeared during
  // insertion of the words in the
  // trie
  Map<Character, Integer> num;

  public TrieNode() {
    this.children = new HashMap<>();
    this.isEndOfWord = false;
    this.num = new HashMap<>();
  }
}

public class Main {
  // Declare root node
  static TrieNode root;

  // Function to create New Trie Node
  static TrieNode getNewTrieNode() {
    TrieNode pNode = new TrieNode();
    pNode.isEndOfWord = false;
    return pNode;
  }

  // Function to insert a string in trie
  static void insertWord(String word) {
    // To hold the value of root
    TrieNode current = root;

    // To hold letters of the word
    char s;

    // Traverse through strings in list
    for (int i = 0; i < word.length(); i++) {
      s = word.charAt(i);

      // If s is not present in the
      // character field of current node
      if (!current.children.containsKey(s)) {
        // Get new node
        TrieNode p = getNewTrieNode();

        // Insert s in character
        // field of current node
        // with reference to node p
        current.children.put(s, p);

        // Insert s in num field
        // of current node with
        // value 1
        current.num.put(s, 1);
      }
      else {
        // Increment the count
        // corresponding to the
        // character s
        current.num.put(s, current.num.get(s) + 1);
      }

      // Go to next node
      current = current.children.get(s);
    }
    current.isEndOfWord = true;
  }

  // Function to count the number of
  // words in trie with given prefix
  static int countWords(Vector<String> words, String prefix) {
    root = getNewTrieNode();

    // Size of list of string
    int n = words.size();

    // Construct trie containing
    // all the words
    for (int i = 0; i < n; i++) {
      insertWord(words.get(i));
    }

    TrieNode current = root;
    char s;

    // Initialize the wordCount = 0
    int wordCount = 0;

    for (int i = 0; i < prefix.length(); i++) {
      s = prefix.charAt(i);

      // If the complete prefix is
      // not present in the trie
      if (!current.children.containsKey(s)) {

        // Make wordCount 0 and
        // break out of loop
        wordCount = 0;
        break;
      }

      // Update the wordCount
      wordCount = current.num.get(s);


      // Go to next node
      current = current.children.get(s);
    }

    return wordCount;
  }


  // Driver Code
  public static void main(String[] args)
  {
    // input list of words
    Vector<String> words = new Vector<>();
    words.add("apk");
    words.add("app");
    words.add("apple");
    words.add("arp");
    words.add("array");

    // Given prefix to find
    String prefix = "ap";

    // Print the number of words with
    // given prefix
    System.out.println(countWords(words, prefix));
  }
}

// This code is contributed by Aman Kumar.

Python3

# Python3 implementation of counting the
# number of words in a trie with a
# given prefix

# Trie Node
class TrieNode:

    def __init__(self):

        # Using map to store the pointers
        # of children nodes for dynamic
        # implementation, for making the
        # program space efficient
        self.children = dict()

        # If isEndOfWord is true, then
        # node represents end of word
        self.isEndOfWord = False

        # num represents number of times
        # a character has appeared during
        # insertion of the words in the
        # trie
        self.num = dict()

# Declare root node
root = None

# Function to create New Trie Node
def getNewTrieNode():

    pNode = TrieNode()
    return pNode

# Function to insert a string in trie
def insertWord(word):

    global root

    # To hold the value of root
    current = root

    # To hold letters of the word
    s = ''

    # Traverse through strings in list
    for i in range(len(word)):
        s = word[i]

        # If s is not present in the
        # character field of current node
        if (s not in current.children):

            # Get new node
            p = getNewTrieNode()

            # Insert s in character
            # field of current node
            # with reference to node p
            current.children[s] = p

            # Insert s in num field
            # of current node with
            # value 1
            current.num[s] = 1
        else:

            # Increment the count
            # corresponding to the
            # character s
            current.num[s] = current.num[s] + 1

        # Go to next node
        current = current.children[s]

    current.isEndOfWord = True

# Function to count the number of
# words in trie with given prefix
def countWords(words, prefix):

    global root
    root = getNewTrieNode()

    # Size of list of string
    n = len(words)

    # Construct trie containing
    # all the words
    for i in range(n):
        insertWord(words[i])

    current = root
    s = ''

    # Initialize the wordCount = 0
    wordCount = 0

    for i in range(len(prefix)):
        s = prefix[i]

        # If the complete prefix is
        # not present in the trie
        if (s not in current.children):

            # Make wordCount 0 and
            # break out of loop
            wordCount = 0
            break

        # Update the wordCount
        wordCount = (current.num)[s]

        # Go to next node
        current = (current.children)[s]

    return wordCount

# Driver Code
if __name__=='__main__':

    # input list of words
    words = [ "apk", "app", "apple",
              "arp", "array" ]

    # Given prefix to find
    prefix = "ap"

    # Print the number of words with
    # given prefix
    print(countWords(words, prefix))

# This code is contributed by rutvik_56

C#

using System;
using System.Collections.Generic;

// Define the TrieNode class to represent a node in the Trie data structure.
public class TrieNode
{
    public Dictionary<char, TrieNode> Children; // Dictionary to store child nodes.
    public bool IsEndOfWord; // Indicates whether this node represents the end of a word.
    public Dictionary<char, int> Num; // Dictionary to store the count of words that start with a particular character.

    // Constructor to initialize the TrieNode object.
    public TrieNode()
    {
        Children = new Dictionary<char, TrieNode>();
        IsEndOfWord = false;
        Num = new Dictionary<char, int>();
    }
}

// Define the Trie class to represent the Trie data structure.
public class Trie
{
    public TrieNode Root; // The root node of the Trie.

    // Constructor to initialize the Trie object.
    public Trie()
    {
        Root = new TrieNode();
    }

    // Method to insert a word into the Trie.
    public void InsertWord(string word)
    {
        TrieNode current = Root;

        foreach (char c in word)
        {
            if (!current.Children.ContainsKey(c)) // If the character is not present in the child nodes, add it.
            {
                current.Children = new TrieNode();
                current.Num = 1; // Since this is the first occurrence of the character, set its count to 1.
            }
            else // If the character is already present in the child nodes, increment its count.
            {
                current.Num++;
            }
            current = current.Children; // Move to the next node.
        }

        current.IsEndOfWord = true; // Set the IsEndOfWord flag to true for the last node in the word.
    }

    // Method to count the number of words in the Trie that start with a given prefix.
    public int CountWords(List<string> words, string prefix)
    {
        Trie trie = new Trie(); // Create a new Trie object.
        foreach (string word in words) // Insert all the words in the input list into the Trie.
        {
            trie.InsertWord(word);
        }

        TrieNode current = trie.Root; // Start at the root node of the Trie.
        int wordCount = 0;

        foreach (char c in prefix) // Traverse the Trie using the characters in the prefix.
        {
            if (!current.Children.ContainsKey(c)) // If the character is not present in the child nodes, there are no words in the Trie that start with the given prefix.
            {
                wordCount = 0;
                break;
            }

            wordCount = current.Num; // Get the count of words that start with the given prefix.
            current = current.Children; // Move to the next node.
        }

        return wordCount; // Return the count of words that start with the given prefix.
    }
}

// Define the Program class to test the Trie data structure.
public class Program
{
    public static void Main()
    {
        List<string> words = new List<string> { "apk", "app", "apple", "arp", "array" };
        string prefix = "ap";
        Trie trie = new Trie();
        int count = trie.CountWords(words, prefix); // Count the number of words that start with the given prefix.
        Console.WriteLine(count); // Print the count.
    }
}
// This code is contributed by rudra1807raj

Javascript

// JavaScript implementation of counting the
// number of words in a trie with a
// given prefix

// Trie Node
class TrieNode {
    constructor() {
        // Using map to store the pointers
        // of children nodes for dynamic
        // implementation, for making the
        // program space efficient
        this.children = new Map();

        // If isEndOfWord is true, then
        // node represents end of word
        this.isEndOfWord = false;

        // num represents number of times
        // a character has appeared during
        // insertion of the words in the
        // trie
        this.num = new Map();
    }
}

// Declare root node
let root = null;

// Function to create New Trie Node
function getNewTrieNode() {
    let pNode = new TrieNode();
    return pNode;
}

// Function to insert a string in trie
function insertWord(word) {
    // To hold the value of root
    let current = root;

    // To hold letters of the word
    let s = '';

    // Traverse through strings in list
    for (let i = 0; i < word.length; i++) {
        s = word[i];

        // If s is not present in the
        // character field of current node
        if (!current.children.has(s)) {
            // Get new node
            let p = getNewTrieNode();

            // Insert s in character
            // field of current node
            // with reference to node p
            current.children.set(s, p);

            // Insert s in num field
            // of current node with
            // value 1
            current.num.set(s, 1);
        } else {
            // Increment the count
            // corresponding to the
            // character s
            current.num.set(s, current.num.get(s) + 1);
        }

        // Go to next node
        current = current.children.get(s);
    }

    current.isEndOfWord = true;
}

// Function to count the number of
// words in trie with given prefix
function countWords(words, prefix) {
    root = getNewTrieNode();

    // Size of list of string
    let n = words.length;

    // Construct trie containing
    // all the words
    for (let i = 0; i < n; i++) {
        insertWord(words[i]);
    }

    let current = root;
    let s = '';

    // Initialize the wordCount = 0
    let wordCount = 0;

    for (let i = 0; i < prefix.length; i++) {
        s = prefix[i];

        // If the complete prefix is
        // not present in the trie
        if (!current.children.has(s)) {
            // Make wordCount 0 and
            // break out of loop
            wordCount = 0;
            break;
        }

        // Update the wordCount
        wordCount = current.num.get(s);

        // Go to next node
        current = current.children.get(s);
    }

    return wordCount;
}

// Driver Code
let words = [ "apk", "app", "apple", "arp", "array" ];
let prefix = "ap";
console.log(countWords(words, prefix));  // Output: 3

Output

3

Time Complexity: O(n*l) where n = the number of words inserted in Trie and l = the length of the longest word inserted in Trie.

Auxiliary Space: O(n*l).