【LeetCode】208. Implement Trie (Prefix Tree)

Implement Trie (Prefix Tree)

Implement a trie with insert, search, and startsWith methods.

Note:
You may assume that all inputs are consist of lowercase letters a-z.

 

一个字母代表一个子树，因此为26叉树，end标记表示是否存在以该字母为结尾的字符串。

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

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

    // Inserts a word into the trie.
    void insert(string word) {
        int i = 0;
        TrieNode* curNode = root;
        while(i < word.size() && curNode->children[word[i]-'a'] != NULL)
        {
            curNode = curNode->children[word[i]-'a'];
            i ++;
        }
        if(i == word.size())
        {
            if(curNode->end == true)
            // exist
                return;
            else
            // insert
                curNode->end = true;
        }
        else
        {
            while(i < word.size())
            {
                curNode->children[word[i]-'a'] = new TrieNode();
                curNode = curNode->children[word[i]-'a'];
                i ++;
            }
            curNode->end = true;
        }
    }

    // Returns if the word is in the trie.
    bool search(string word) {
        int i = 0;
        TrieNode* curNode = root;
        while(i < word.size() && curNode->children[word[i]-'a'] != NULL)
        {
            curNode = curNode->children[word[i]-'a'];
            i ++;
        }
        if(i == word.size() && curNode->end == true)
        // curNode must be leaf
            return true;
        else
            return false;
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    bool startsWith(string prefix) {
        int i = 0;
        TrieNode* curNode = root;
        while(i < prefix.size() && curNode->children[prefix[i]-'a'] != NULL)
        {
            curNode = curNode->children[prefix[i]-'a'];
            i ++;
        }
        if(i == prefix.size())
        // curNode might be left or not
            return true;
        else
            return false;
    }

private:
    TrieNode* root;
};

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