208. Implement Trie (Prefix Tree)

Implement a trie with insert, search, and startsWith methods.

Note:

You may assume that all inputs are consist of lowercase letters a-z.

class TrieNode{
public:
    char content;
    bool is_end;
    int shared;
    vector<TrieNode *> children;
public:
    TrieNode(char cnt):content(cnt),is_end(false),shared(0)
    {
        
    }
    TrieNode *sub(char ch)
    {
        if(!children.empty())
        {
            vector<TrieNode *>::iterator it = children.begin();
            for(it ; it!=children.end(); it++)
            {
                if((*it)->content == ch)
                    return *it;
            }
        }
        return NULL;
    }
};

class Trie {
public:
    /** Initialize your data structure here. */
    Trie() {
        root = new TrieNode(' ');
    }
    
    /** Inserts a word into the trie. */
    void insert(string word) {
       TrieNode *p = root;
       for(int i=0; i<word.size(); i++)
       {
           if(p->sub(word[i])!=NULL)
           {
               p = p->sub(word[i]);
           }else
           {
               TrieNode *node = new TrieNode(word[i]);
               p->children.push_back(node);
               p = node;
           }
       }
       p->is_end = true;
    }
    
    /** Returns if the word is in the trie. */
    bool search(string word) {
       TrieNode *p = root;
       for(int i=0; i<word.size(); i++)
       {
           if(p->sub(word[i])==NULL)
           {
               return false;
           }else
           {
              p = p->sub(word[i]);
           }
       }
       return p->is_end == true;
    }
    
    /** Returns if there is any word in the trie that starts with the given prefix. */
    bool startsWith(string prefix) {
       TrieNode *p = root;
       for(int i=0; i<prefix.size(); i++)
       {
           if(p->sub(prefix[i])==NULL)
           {
               return false;
           }else
           {
              p = p->sub(prefix[i]);
           }
       }
       return true;
    }
private:
    TrieNode *root;
};