LeetCode: Implement Trie (Prefix Tree)

Implement Trie (Prefix Tree)

#include <iostream>
using namespace std;

class TrieNode {
public:
    // Initialize your data structure here.
    TrieNode() {
        for(int i=0; i<26; i++){
            next[i] =  NULL;
        }
        complete = false;
    }
    TrieNode* next[26];
    bool complete;
};

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

    // Inserts a word into the trie.
    void insert(string word) {
        TrieNode *p = root;
        int len = word.size();
        for(int i=0; i<len; i++){
            int c = word[i]-'a';
            if(p->next[c] == NULL){
                p->next[c] = new TrieNode();
            }
            p = p->next[c];
        }
        p->complete = true;
    }

    // Returns if the word is in the trie.
    bool search(string word) {
        int len = word.size();
        TrieNode* p = root;
        for(int i=0; i<len; i++){
            int c = word[i]-'a';
            if(p->next[c] == NULL){
                return false;
            }
            p = p->next[c];
        }

        return p->complete;
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    bool startsWith(string prefix) {
        int len = prefix.size();
        TrieNode* p = root;
        for(int i=0; i<len; i++){
            int c = prefix[i]-'a';
            if(p->next[c] == NULL){
                return false;
            }
            p = p->next[c];
        }
        return true;
    }

private:
    TrieNode* root;
};

int main()
{
    Trie tmp;
    tmp.insert("abc");
    tmp.insert("az");
    cout<<tmp.startsWith("ab")<<endl;
    cout<<tmp.startsWith("ad")<<endl;
    cout<<tmp.search("ab")<<endl;
    cout<<tmp.search("abc")<<endl;

    return 0;
}