Implement Trie (Prefix Tree) Solutions in C++

// Source : https://leetcode.com/problems/implement-trie-prefix-tree/
// Author : Hao Chen
// Date   : 2015-06-09




class TrieNode {
public:
    // Initialize your data structure here.
    TrieNode():isWord(false) {   }
    unordered_map<char, TrieNode*> children;
    bool isWord;
};

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

    // Inserts a word into the trie.
    void insert(string s) {
        if (s.size()<=0) return;
        TrieNode * node = root;
        for (int i=0; i<s.size(); i++) {
            if (node->children.find(s[i]) == node->children.end()){
                node->children[s[i]] = new TrieNode();
            }
            node = node->children[s[i]];
        }
        node->isWord = true;
    }

    // Returns if the word is in the trie.
    bool search(string key) {
        return retrieve(key, true);
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    bool startsWith(string prefix) {
        return retrieve(prefix, false);
    }

private:

    inline bool retrieve(const string& key, bool isWord) {
        if (key.size() <= 0) return false;
        TrieNode * node = root;
        for (int i=0; i<key.length(); i++) {
            if (node->children.find(key[i]) == node->children.end()) {
                return false;
            }
            node = node->children[key[i]];
        }
        return isWord ? node->isWord : true;
    }

    TrieNode* root;
};

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

/// Source : https://leetcode.com/problems/implement-trie-prefix-tree/description/
/// Author : liuyubobobo
/// Time   : 2017-11-05

#include <iostream>
#include <vector>
#include <map>

using namespace std;

/// Trie Recursive version
class Trie{

private:
    struct Node{
        map<char, int> next;
        bool end = false;
    };
    vector<Node> trie;

public:
    Trie(){
        trie.clear();
        trie.push_back(Node());
    }

    /** Inserts a word into the trie. */
    void insert(const string& word){
        insert(0, word, 0);
    }

    /** Returns if the word is in the trie. */
    bool search(const string& word){
        return search(0, word, 0);
    }

    /** Returns if there is any word in the trie that starts with the given prefix. */
    bool startsWith(const string& prefix) {
        return startsWith(0, prefix, 0);
    }

private:
    void insert(int treeID, const string& word, int index){

        if(index == word.size()) {
            trie[treeID].end = true;
            return;
        }

        if(trie[treeID].next.find(word[index]) == trie[treeID].next.end()){
            trie[treeID].next[word[index]] = trie.size();
            trie.push_back(Node());
        }

        insert(trie[treeID].next[word[index]], word, index + 1);
    }

    bool search(int treeID, const string& word, int index){

        if(index == word.size())
            return trie[treeID].end;

        if(trie[treeID].next.find(word[index]) == trie[treeID].next.end())
            return false;

        return search(trie[treeID].next[word[index]], word, index + 1);
    }

    bool startsWith(int treeID, const string& prefix, int index){

        if(index == prefix.size())
            return true;

        if(trie[treeID].next.find(prefix[index]) == trie[treeID].next.end())
            return false;

        return startsWith(trie[treeID].next[prefix[index]], prefix, index + 1);
    }
};


void printBool(bool res){
    cout << (res ? "True" : "False") << endl;
}

int main() {

    Trie trie1;
    trie1.insert("ab");
    printBool(trie1.search("a"));     // false
    printBool(trie1.startsWith("a")); // true;

    cout << endl;

    // ---

    Trie trie2;
    trie2.insert("a");
    printBool(trie2.search("a"));     // true
    printBool(trie2.startsWith("a")); // true;

    return 0;
}

/// Source : https://leetcode.com/problems/implement-trie-prefix-tree/description/
/// Author : liuyubobobo
/// Time   : 2017-11-05

#include <iostream>
#include <vector>
#include <map>

using namespace std;

/// Trie Non-Recursive version
class Trie{

private:
    struct Node{
        map<char, int> next;
        bool end = false;
    };
    vector<Node> trie;

public:
    Trie(){
        trie.clear();
        trie.push_back(Node());
    }

    /** Inserts a word into the trie. */
    void insert(const string& word){

        int treeID = 0;
        for(char c: word){
            if(trie[treeID].next.find(c) == trie[treeID].next.end()){
                trie[treeID].next[c] = trie.size();
                trie.push_back(Node());
            }

            treeID = trie[treeID].next[c];
        }

        trie[treeID].end = true;
    }

    /** Returns if the word is in the trie. */
    bool search(const string& word){

        int treeID = 0;
        for(char c: word){
            if(trie[treeID].next.find(c) == trie[treeID].next.end())
                return false;

            treeID = trie[treeID].next[c];
        }
        return trie[treeID].end;
    }

    /** Returns if there is any word in the trie that starts with the given prefix. */
    bool startsWith(const string& prefix) {

        int treeID = 0;
        for(char c: prefix){
            if(trie[treeID].next.find(c) == trie[treeID].next.end())
                return false;

            treeID = trie[treeID].next[c];
        }
        return true;
    }
};


void printBool(bool res){
    cout << (res ? "True" : "False") << endl;
}

int main() {

    Trie trie1;
    trie1.insert("ab");
    printBool(trie1.search("a"));     // false
    printBool(trie1.startsWith("a")); // true;

    cout << endl;

    // ---

    Trie trie2;
    trie2.insert("a");
    printBool(trie2.search("a"));     // true
    printBool(trie2.startsWith("a")); // true;

    return 0;
}