Solution Review: Search Auto-Complete Optimized

#include <iostream>
#include <vector>
#include <set>
using namespace std;

// Definition of a trie node
class TrieNode
{
    public:
       	// children array to store pointer to TrieNode of next 26 alphabets
        TrieNode *children[26];
        // a set to store all the words passing through this trie node
        set<string> wordsPassingFromHere;
};

class Trie
{
    TrieNode * root;

    public:
        Trie()
        {
            root = new TrieNode();
        }

        // Inserts the string in trie
        void insert(string s)
        {
            // Points curr to the root of trie.
            TrieNode *curr = root;

            // iterate through the chacraters of the word
            for (char &c: s)
            {
                // if a child node is not found, create one
                if (curr->children[c - 'a'] == nullptr)
                    curr->children[c - 'a'] = new TrieNode();

                // move to the child node
                curr = curr->children[c - 'a'];

                // insert the word in current node's wordsPassingFromHere list
                curr->wordsPassingFromHere.insert(s);

                // if after insertion, the set size exceeds 3
                // erase the last element of the set
                if (curr->wordsPassingFromHere.size() > 3)
                {
                    // Get the iterator
                    set<string>::iterator it;

                    // Get the position of last element to be deleted
                    it = prev(curr->wordsPassingFromHere.end());
                    
                    // erase the last element from set
                    curr->wordsPassingFromHere.erase(it);
                }
            }
        }

        vector<string> returnSuggestions(set<string> suggestions){
            // store them in a vector and return 
            vector<string> v(suggestions.begin(), suggestions.end());
            return v;
        }
        
        vector<string> getWordsStartingWith(string & prefix)
        {
            // point the current node to the root
            TrieNode *curr = root;

            // Move curr to the end of prefix in its trie representation
            for (char &c: prefix)
            {
                // if the complete query word is not found 
                // in the store, terminate the search and return results
                if (curr->children[c - 'a']==nullptr)
                    return returnSuggestions(curr->wordsPassingFromHere);
                
                // move to the child node
                curr = curr->children[c - 'a'];
            }

            //retrieve the words to be returned
            return returnSuggestions(curr->wordsPassingFromHere);
        }
};

vector<vector < string>> autoSuggest(vector<string> &words, string searchWord)
{
   	// initialize a trie
    Trie *trie = new Trie();

   	// initialize a result vector
    vector<vector < string>> result;

   	// insert all words into trie
    for (string w: words)
        trie->insert(w);

    string prefix;

   	// iterate through word's characters
    for (char c: searchWord)
    {
       	// keep appending letters to the prefix
        prefix += c;
       	// fetch words starting with the current prefix
        result.push_back(trie->getWordsStartingWith(prefix));
    }

    return result;
}

// Definition of a trie node
class TrieNode
{
    public:
       	// children array to store pointer to TrieNode of next 26 alphabets
        TrieNode *children[26];
        // a set to store all the words passing through this trie node
        set<string> wordsPassingFromHere;
};