Implement Trie II (Prefix Tree)

Link

Description

A trie (pronounced as "try") or prefix tree is a tree data structure used to efficiently store and retrieve keys in a dataset of strings. There are various applications of this data structure, such as autocomplete and spellchecker.

Implement the Trie class:

• Trie() Initializes the trie object.
• void insert(String word) Inserts the string word into the trie.
• int countWordsEqualTo(String word) Returns the number of instances of the string word in the trie.
• int countWordsStartingWith(String prefix) Returns the number of strings in the trie that have the string prefix as a prefix.
• void erase(String word) Erases the string word from the trie.

Example 1:

• Input:

["Trie", "insert", "insert", "countWordsEqualTo", "countWordsStartingWith", "erase", "countWordsEqualTo", "countWordsStartingWith", "erase", "countWordsStartingWith"]
[[], ["apple"], ["apple"], ["apple"], ["app"], ["apple"], ["apple"], ["app"], ["apple"], ["app"]]

• Output:

[null, null, null, 2, 2, null, 1, 1, null, 0]

• Explanation:

Trie trie = new Trie();
trie.insert("apple");               // Inserts "apple".
trie.insert("apple");               // Inserts another "apple".
trie.countWordsEqualTo("apple");    // There are two instances of "apple" so return 2.
trie.countWordsStartingWith("app"); // "app" is a prefix of "apple" so return 2.
trie.erase("apple");                // Erases one "apple".
trie.countWordsEqualTo("apple");    // Now there is only one instance of "apple" so return 1.
trie.countWordsStartingWith("app"); // return 1
trie.erase("apple");                // Erases "apple". Now the trie is empty.
trie.countWordsStartingWith("app"); // return 0

Constraints:

• 1 <= word.length, prefix.length <= 2000
• word and prefix consist only of lowercase English letters.
• At most 3 * 10^4 calls in total will be made to insert, countWordsEqualTo, countWordsStartingWith, and erase.
• It is guaranteed that for any function call to erase, the string word will exist in the trie.

Solution

Original Solution

struct TrieNode {
    TrieNode* nexts[26];
    int passing = 0;
    int ending = 0;

    ~TrieNode() {
        for (auto& node : nexts) {
            delete node;
        }
    }
};

class Trie {
public:
    TrieNode* root;

    Trie() {
        root = new TrieNode();
    }

    ~Trie() {
        delete root;
    }

    void insert(string word) {
        TrieNode* node = root;

        for (const char& ch : word) {
            int index = ch - 'a';
            if (!node->nexts[index]) {
                node->nexts[index] = new TrieNode();
            }

            node = node->nexts[index];
            node->passing++;
        }

        node->ending++;
    }

    int countWordsEqualTo(string word) {
        TrieNode* node = root;

        for (const char& ch : word) {
            int index = ch - 'a';
            if (!node->nexts[index]) {
                return 0;
            }

            node = node->nexts[index];
        }

        return node->ending;
    }

    int countWordsStartingWith(string prefix) {
        TrieNode* node = root;

        for (const char& ch : prefix) {
            int index = ch - 'a';
            if (!node->nexts[index]) {
                return 0;
            }

            node = node->nexts[index];
        }

        return node->passing;
    }

    void erase(string word) {
        TrieNode *node = root;

        for (const char& ch : word) {
            int index = ch - 'a';
            node = node->nexts[index];
            node->passing--;
        }

        node->ending--;
    }
};

/**
 * Your Trie object will be instantiated and called as such:
 * Trie* obj = new Trie();
 * obj->insert(word);
 * int param_2 = obj->countWordsEqualTo(word);
 * int param_3 = obj->countWordsStartingWith(prefix);
 * obj->erase(word);
 */

Time and space complexity of each function:

1. Constructor (Trie()):
2. Time Complexity: $O(1)$. The constructor only initializes the root, which takes constant time.

3. Space Complexity: $O(1)$ for the allocation of the root node.

4. Insert (insert(string word)):

5. Time Complexity: $O(L)$, where $L$ is the length of the word being inserted. This is because it traverses the Trie for the length of the word.

6. Space Complexity: $O(L)$ in the worst case when a new path is created for the entire length of the word. This accounts for the space used by the new nodes.

7. Count Words Equal To (countWordsEqualTo(string word)):

8. Time Complexity: $O(L)$, where $L$ is the length of the word. It traverses the Trie for the length of the word to find the ending count.

9. Space Complexity: $O(1)$, as it only uses a pointer to traverse the nodes.

10. Count Words Starting With (countWordsStartingWith(string prefix)):

11. Time Complexity: $O(P)$, where $P$ is the length of the prefix. The method traverses the Trie for the length of the prefix.

12. Space Complexity: $O(1)$, since it similarly uses a pointer for traversal.

13. Erase (erase(string word)):

14. Time Complexity: $O(L)$, where $L$ is the length of the word. It decreases the passing count while traversing the Trie.
15. Space Complexity: $O(1)$, as it utilizes a pointer to navigate the Trie.

More general way of writing erase function:

void erase(string word) {
        if (countWordsEqualTo(word) <= 0) return;  // unnecessary for this question

        TrieNode *node = root;

        for (const char& ch : word) {
            int index = ch - 'a';
            if (--node->nexts[index]->passing == 0) {
                node->nexts[index] = nullptr;
                delete node->nexts[index];
                return;
            }
            node = node->nexts[index];
        }

        node->ending--;
    }

Using Hash Map

struct TrieNode {
    unordered_map<char, TrieNode*> nexts;
    int passing = 0;
    int ending = 0;

    ~TrieNode() {
        for (auto& pair : nexts) {
            delete pair.second;
        }
    }
};

class Trie {
public:
    TrieNode* root;

    Trie() : root(new TrieNode()) {}

    ~Trie() {
        delete root;
    }

    void insert(string word) {
        auto node = root;
        node->passing++;

        for (const auto& ch : word) {
            if (!node->nexts.count(ch)) {
                node->nexts[ch] = new TrieNode();
            }

            node = node->nexts[ch];
            node->passing++;
        }

        node->ending++;
    }

    int countWordsEqualTo(string word) {
        auto node = root;

        for (const auto& ch : word) {
            if (!node->nexts.count(ch)) {
                return 0;
            }

            node = node->nexts[ch];
        }

        return node->ending;
    }

    int countWordsStartingWith(string prefix) {
        auto node = root;

        for (const auto& ch : prefix) {
            if (!node->nexts.count(ch)) {
                return 0;
            }

            node = node->nexts[ch];
        }

        return node->passing;
    }

    void erase(string word) {
        auto node = root;
        node->passing--;

        for (const auto& ch : word) {
            node = node->nexts[ch];
            node->passing--;
        }

        node->ending--;
    }
};

/**
 * Your Trie object will be instantiated and called as such:
 * Trie* obj = new Trie();
 * obj->insert(word);
 * int param_2 = obj->countWordsEqualTo(word);
 * int param_3 = obj->countWordsStartingWith(prefix);
 * obj->erase(word);
 */

Static Array Solution

class Trie {
private:
    const static int N = 31010;
    int son[N][26] {0};
    int passing[N] {0};
    int ending[N] {0};
    int cnt = 0;

public:
    Trie() {

    }

    void insert(string word) {
        int current = 0;
        passing[current]++;

        for (char& ch : word) {
            int index = ch - 'a';
            if (!son[current][index]) {
                son[current][index] = ++cnt;
            }

            current = son[current][index];
            passing[current]++;
        }

        ending[current]++;
    }

    int countWordsEqualTo(string word) {
        int current = 0;

        for (char& ch : word) {
            int index = ch - 'a';
            if (!son[current][index]) {
                return 0;
            }

            current = son[current][index];
        }

        return ending[current];
    }

    int countWordsStartingWith(string prefix) {
        int current = 0;

        for (char& ch : prefix) {
            int index = ch - 'a';
            if (!son[current][index]) {
                return 0;
            }

            current = son[current][index];
        }

        return passing[current];
    }

    void erase(string word) {
        int current = 0;
        passing[current]--;

        for (char& ch : word) {
            int index = ch - 'a';
            current = son[current][index];
            passing[current]--;
        }

        ending[current]--;
    }
};

/**
 * Your Trie object will be instantiated and called as such:
 * Trie* obj = new Trie();
 * obj->insert(word);
 * int param_2 = obj->countWordsEqualTo(word);
 * int param_3 = obj->countWordsStartingWith(prefix);
 * obj->erase(word);
 */