Friday, January 22, 2016

Trie Implementation

class TrieNode {
public:
    TrieNode* dict[26];
    bool hasWord;   
    // Initialize your data structure here.
    TrieNode() {
        memset(dict,0,sizeof(dict));
        hasWord = false;
    }
};

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

// Inserts a word into the trie.
void insert(string word) {
    int i = 0;
    TrieNode * curr = root;
    while(i
        if( curr->dict[word[i] - 'a'] == nullptr){
            curr->dict[word[i]-'a'] = new TrieNode();
        }
        if(i == word.size()-1){
        curr->dict[word[i]-'a']->hasWord = true;
        }
        curr = curr->dict[word[i] - 'a'];
        i++;
    }
    return;
}

// Returns if the word is in the trie.
    bool search(string word) {
        int i = 0;
        TrieNode *curr = root;
        while(i            if(curr->dict[word[i]-'a'] == nullptr){
                return false;
            }
            if( i == word.size()-1){
                return curr->dict[word[i]-'a']->hasWord;
            }
            curr = curr->dict[word[i] - 'a'];
            i++;
        }
        return false;       
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    bool startsWith(string prefix) {
        TrieNode *curr = root;
        for(auto c:prefix){
            if(curr->dict[c-'a'] == nullptr)
                return false;
            curr = curr->dict[c-'a'];   
        }
        return true;
    }

private:
    TrieNode* root;
};

No comments: