Implement Trie (Prefix Tree)

Problem

A prefix tree (also known as a trie) is a tree data structure used to efficiently store and retrieve keys in a set of strings. Some applications of this data structure include auto-complete and spell checker systems.

Implement the PrefixTree class:

• PrefixTree() Initializes the prefix tree object.
• void insert(String word) Inserts the string word into the prefix tree.
• boolean search(String word) Returns true if the string word is in the prefix tree (i.e., was inserted before), and false otherwise.
• boolean startsWith(String prefix) Returns true if there is a previously inserted string word that has the prefix prefix, and false otherwise.

Examples

Example 1:

Input:
["Trie", "insert", "dog", "search", "dog", "search", "do", "startsWith", "do", "insert", "do", "search", "do"]

Output:
[null, null, true, false, true, null, true]

Explanation:
PrefixTree prefixTree = new PrefixTree();
prefixTree.insert("dog");
prefixTree.search("dog");    // return true
prefixTree.search("do");     // return false
prefixTree.startsWith("do"); // return true
prefixTree.insert("do");
prefixTree.search("do");     // return true

Constraints

• 1 <= word.length, prefix.length <= 1000
• word and prefix are made up of lowercase English letters.

You should aim for a solution with O(n) time for each function call and O(t) space, where n is the length of the given string and t is the total number of nodes created in the Trie.

Solution

Tree where node has map of char to child, and a flag if its the end of a word. To insert a node, we walk down the word and tree in sync, creating any node that doesn’t exist, and set flag at the end. For search, we walk the tree and return false if early breaking, or end does not have flag set. Prefix is similar, but always return true at the end.

struct TrieNode {
    std::unordered_map<char, TrieNode*> children;
    bool end_of_word = false;
};
 
class PrefixTree {
    TrieNode *root;
public:
    PrefixTree() {
        root = new TrieNode();
    }
    ~PrefixTree() { delete root; }
    
    void insert(const string &word) {
        // Start at root, and walk down in step with word
        // If no node exists for this char, make child
        TrieNode *current = root;
        for (const auto &c : word) {
            if (!current->children.contains(c)) {
                current->children[c] = new TrieNode();
            }
            current = current->children[c];
        }
        current->end_of_word = true;
    }
    
    bool search(const string &word) {
        // Start at root, and walk down in step with word
        // If child doesn't contain corresponding char, word not found
        TrieNode *current = root;
        for (const auto &c : word) {
            if (!current->children.contains(c)) {
                return false;
            }
            current = current->children[c];
        }
        // Otherwise check if we are actually at the end of a word
        return current->end_of_word;
    }
    
    bool startsWith(const string &prefix) {
        TrieNode *current = root;
        for (const auto &c : prefix) {
            if (!current->children.contains(c)) {
                return false;
            }
            current = current->children[c];
        }
        // Either we found word or another word continues on
        // Either way, true
        return true;
    }
};