LeetCode: Implement Trie (Prefix Tree)

Problem statement

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

Example:

Trie trie = new Trie();

trie.insert("apple");
trie.search("apple");   // returns true
trie.search("app");     // returns false
trie.startsWith("app"); // returns true
trie.insert("app");   
trie.search("app");     // returns true

Note:

• You may assume that all inputs are consist of lowercase letters a-z.
• All inputs are guaranteed to be non-empty strings.

Solution Thought Process

Each of the trie node has two elements.

• isString , if this node marks an end to a string
• unordered_map<char, TrieNode*>, a hash map of characters which map to TrieNodes.

First, we declare the structure of the trie node and the root of trie globally.

struct TrieNode {
    bool isString = false;
    unordered_map<char,TrieNode*> leaves;
};

TrieNode* root = new TrieNode();

How does insertion work?

Suppose we have to words "abc" and "abd". The first word is already stored in the Trie. The Trie looks like this:

root ---->  a (isString = false)  ---->  b (isString = false) ----> c (isString = true)

If we look at the string "abd", we can see the first two characters are available in the Trie. The second node has c as leave, but it doesn't have d as a leave. So we add it. But it marks the end of the string, so we make the isString of that node to be true.

root ---->  a (isString = false)  ---->  b (isString = false) ----> c (isString = true) 
                                         | 
                                         |
                                         -------------> d (isString = true)

How does the search work?

We start from the root and search for the characters of the string one by one.

• For "abd" in the previous example, first, we see if the leaves of the root and check for 'a', we have found it. We go to the node 'a'.
• From node 'a', we search the leaves to find 'b', we have found it and go to node 'b'.
• From node 'b', we search the leaves to find 'd', we have found it and go to node 'd'.
• In the node 'd', we look at the isString boolean. If it's true, that means that we have found one string ends on this node. Because we have iterated the whole string and came to this node, we can say that the string is available in the trie.

How does prefix search work?

Suppose we have to search for "ab" prefix in the trie in the above example. We start from the root and search for the characters of the string one by one.

• First, we see if the leaves of the root and check for 'a', we have found it. We go to the node 'a'.
• From node 'a', we search the leaves to find 'b', we have found it and go to node 'b'.

We have found both characters in that order in the trie, which means it exists as a prefix. The solution is similar to the search, but it doesn't care about isString.

Solution

class Trie {
private:
    struct TrieNode {
        bool isString = false;
        unordered_map<char,TrieNode*> leaves;
    };

    TrieNode* root = new TrieNode();
public:
    /** Initialize your data structure here. */
    Trie() {

    }

    /** Inserts a word into the trie. */
    void insert(string word) {
        TrieNode* p = root;
        for(char c: word)
        {
            if(p->leaves.find(c)==p->leaves.end())
            {
                p->leaves[c] = new TrieNode();
            }
            p = p->leaves[c];
        }

        if(p->isString==false)
        {
            p->isString = true;
        }
    }

    /** Returns if the word is in the trie. */
    bool search(string word) {
        TrieNode* p = root;
        for(char c:word)
        {
            if(p->leaves.find(c)==p->leaves.end())
            {
                return false;
            }
            p = p->leaves[c];
        }
        if(p->isString)
        {
            return true;
        }
        return false;
    }

    /** Returns if there is any word in the trie that starts with the given prefix. */
    bool startsWith(string prefix) {
        TrieNode* p = root;
        for(char c:prefix)
        {
            if(p->leaves.find(c) == p->leaves.end())
            {
                return false;
            }
            p = p->leaves[c];
        }
        return true;

    }
};

/**
 * Your Trie object will be instantiated and called as such:
 * Trie* obj = new Trie();
 * obj->insert(word);
 * bool param_2 = obj->search(word);
 * bool param_3 = obj->startsWith(prefix);
 */

Complexity

Time Complexity :

For inserting, O(nk), k = is the average word length, n = number of words

For Searching, O(k), where k is the search key length.