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:
TrieNode *nodeMap[26];
int endNum;
public:
// Initialize your data structure here.
TrieNode() {
for (int i=0; i<26; i++)
nodeMap[i] = NULL;
endNum = 0;
}
};
class Trie {
public:
Trie() {
root = new TrieNode();
}
// Inserts a word into the trie.
void insert(string word) {
TrieNode *curNode = root;
for (int i=0; i<word.size(); i++)
{
int wordIdx = word[i]-'a';
if (curNode->nodeMap[wordIdx] == NULL)
{
TrieNode *newNode = new TrieNode();
curNode->nodeMap[wordIdx] = newNode;
}
curNode = curNode->nodeMap[wordIdx];
}
curNode->endNum++;
}
// Returns if the word is in the trie.
bool search(string word) {
TrieNode *curNode = root;
for (int i=0; i<word.size(); i++)
{
int wordIdx = word[i]-'a';
if (curNode->nodeMap[wordIdx] == NULL)
return false;
curNode = curNode->nodeMap[wordIdx];
}
if (curNode->endNum)
return true;
return false;
}
// Returns if there is any word in the trie
// that starts with the given prefix.
bool startsWith(string prefix) {
TrieNode *curNode = root;
for (int i=0; i<prefix.size(); i++)
{
int wordIdx = prefix[i]-'a';
if (curNode->nodeMap[wordIdx] == NULL)
return false;
curNode = curNode->nodeMap[wordIdx];
}
return true;
}
private:
TrieNode* root;
};
// Your Trie object will be instantiated and called as such:
// Trie trie;
// trie.insert("somestring");
// trie.search("key");