leecode_208 Implement Trie (Prefix Tree)

Implement a trie with 

insert

, 

search

,
and 

startsWith

 methods.

class TrieNode {
public:
// Initialize your data structure here.
char c;
unordered_map<char, TrieNode*> sub;

TrieNode() {
c='0';
}

TrieNode(char parent) {
c=parent;
}

void add_sub(char c) {
TrieNode* node=new TrieNode(c);
sub[c]=node;
}

bool has_sub(char c)
{
if (sub.find(c)!=sub.end())
return true;
else
return false;
}

TrieNode* get_sub(char c){
return sub[c];
}

bool excited(){
if (sub.find('0')!=sub.end())
return true;
else
return false;
}

};

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

// Inserts a word into the trie.
void insert(string word) {
TrieNode* head=root;
TrieNode* node;

for (int i=0; i<word.size();i++){
if (!head->has_sub(word[i]))
head->add_sub(word[i]);

node=head->get_sub(word[i]);
head=node;
}

head->add_sub('0');
}

// Returns if the word is in the trie.
bool search(string word) {
TrieNode* head=root;
TrieNode* node;

for (int i=0;i<word.size();i++){
if (!head->has_sub(word[i]))
return false;
node=head->get_sub(word[i]);
head=node;
//cout<<"find "<<word[i]<<endl;
}

return (*node).excited();

}

// Returns if there is any word in the trie
// that starts with the given prefix.
bool startsWith(string prefix) {
TrieNode* head=root;
TrieNode* node;

for (int i=0;i<prefix.size();i++){
if (!head->has_sub(prefix[i]))
return false;
node=head->get_sub(prefix[i]);
head=node;
}

for (int i=0;i<26;i++){
if(head->has_sub('a'+i))
return true;
}

return head->excited();
}

private:
TrieNode* root;
};

// Your Trie object will be instantiated and called as such:
// Trie trie;
// trie.insert("somestring");
// trie.search("key");