Skip to content
Related Articles

Related Articles

Count the number of words with given prefix using Trie

View Discussion
Improve Article
Save Article
Like Article
  • Difficulty Level : Easy
  • Last Updated : 31 Jan, 2022

Prerequisite: Trie 

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

Examples: 

Input: str = [ “apk”, “app”, “apple”, “arp”, “array” ], pre = “ap” 
Output:
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:

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 end of word node
  3. num: This field is used to count the number of times a node is visited during insertion in trie

Steps: 

  • Insert 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.
  • Value of 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;
}

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

Output: 

3

 

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


My Personal Notes arrow_drop_up
Recommended Articles
Page :

Start Your Coding Journey Now!