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 list of string with 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 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:

Below is the implementation of the above approach:

filter_none

edit
close

play_arrow

link
brightness_4
code

// 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 efiicient
    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
            // corressponding 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;
}
chevron_right

Output:
3

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





Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.


Article Tags :
Practice Tags :