Here is the implementation of Trie Tree. Using array is good but using hashmap is more straightforward.
The key interface need to be implement is insert and search.
#include <iostream>
#include <vector>
#include <unordered_map>
using namespace std;
class TrieNode {
public:
unordered_map<char, TrieNode*> children;
bool isWord;
TrieNode(){
isWord = false;
}
void insert(string word, int index){
if(index == word.length()){
this->isWord = true;
return;
}
char ch = word.at(index);
if(children.find(ch) == children.end()){
children[ch] = new TrieNode();
}
children[ch]->insert(word, index+1);
}
// recursive
TrieNode* find(string word, int index){
if(index == word.length()){
return this;
}
char ch = word.at(index);
if(children.find(ch) == children.end()){
return NULL;
}
return children[ch]->find(word, index+1);
}
// non recursive
TrieNode * find(string word){
if(word.length() == 0){
return NULL;
}
TrieNode* node = NULL;
unordered_map<char, TrieNode*> sons = children;
for(int i = 0; i < word.length(); i++){
char ch = word.at(i);
if(sons.find(ch) == sons.end()){
return NULL;
}
node = sons[ch];
sons = node->children;
}
return node;
}
};
class Trie {
private:
TrieNode * root;
public:
Trie(){
root = new TrieNode();
}
Trie(vector<string> dictionary){
root = new TrieNode();
// build the trie tree
for(string word : dictionary){
root->insert(word, 0);
}
}
void insert(string word){
root->insert(word, 0);
}
bool search(string word){
TrieNode* node = root->find(word, 0);
return (node != NULL) && (node->isWord);
}
bool startWith(string prefix){
TrieNode* node = root->find(prefix, 0);
return node != NULL;
}
///TEST PURPOSE
// get all strings that start with prefix
void helper(vector<string> &result, string &path, TrieNode * node, string prefix){
if(node == NULL || node->children.size() == 0){
result.push_back(prefix + path);
return;
}
for(auto it = node->children.begin();
it != node->children.end();
it++){
path.push_back(it->first);
helper(result, path, it->second, prefix);
path.pop_back();
}
}
vector<string> getAll(string prefix){
vector<string> ans;
TrieNode * node = root->find(prefix, 0);
if(node == NULL || node->children.size() == 0){
return ans;
}
string path;
helper(ans, path, node, prefix);
return ans;
}
};
int main(void){
vector<string> dictionary = {"hello world", "goodman", "good engineer", "emc", "sql language", "google", "amazon"};
Trie trie(dictionary);
vector<string> output = trie.getAll("goo");
for(string str : output){
cout << str <<endl;
}
return 0;
}