Thursday, May 14, 2015

[LeetCode] Implement Trie (Prefix Tree)

Implement a trie with insertsearch, and startsWith methods.
Note:
You may assume that all inputs are consist of lowercase letters a-z.
Analysis: There is no need to keep an another class called TrieNode in this case.

class Trie {
public:
    // Initialize your data structure here.
    void insert(string s) {
        Trie *cur = this;
        for(int i=0;i<s.length();i++)
        {
            int index = s[i]-'a';
            if(cur->current[index]==NULL){
                cur->current[index] = new Trie();
            }
            cur = cur->current[index];
        }
        cur->endOfWord = true;
    }

    bool startsWith(string prefix) {
        Trie *cur = this;
        for(int i=0;i<prefix.length();i++)
        {
            cur = cur->current[prefix[i]-'a'];
            if(!cur) return false;
        }
        return true;
    }
    
    bool search(string key)
    {
        Trie *cur = this;
        for(int i=0;i<key.length();i++)
        {
            cur = cur->current[key[i] - 'a'];
            if(!cur) return false;
        }
        return cur->endOfWord;
    }
    
    Trie() {
        endOfWord = false;
        for(int i=0;i<26;i++) current[i]=NULL;
    }
    
    ~Trie(){
        for(int i=0;i<26;i++)
        {
            if(current[i]) delete current[i];
        }
    }
private:
    Trie *current[26];
    bool endOfWord;
};