1804. Implement Trie II (Prefix Tree)

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 * 104 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.

1804. Implement Trie II (Prefix Tree)
use std::collections::HashMap;

#[derive(Default)]
struct Trie {
    children: HashMap<char, Trie>,
    count: i32,
}

impl Trie {
    fn new() -> Self {
        Trie::default()
    }

    fn insert(&mut self, word: String) {
        let mut link = self;
        for c in word.chars() {
            link = link.children.entry(c).or_default();
        }
        link.count += 1;
    }

    fn count_words_equal_to(&self, word: String) -> i32 {
        let mut link = self;
        for c in word.chars() {
            if let Some(next) = link.children.get(&c) {
                link = next;
            } else {
                return 0;
            }
        }
        link.count
    }

    fn count_words_starting_with(&self, prefix: String) -> i32 {
        let mut link = self;
        for c in prefix.chars() {
            if let Some(next) = link.children.get(&c) {
                link = next;
            } else {
                return 0;
            }
        }
        link.dfs()
    }

    fn dfs(&self) -> i32 {
        let mut res = self.count;
        for next in self.children.values() {
            res += next.dfs();
        }
        res
    }

    fn erase(&mut self, word: String) {
        let mut link = self;
        for c in word.chars() {
            link = link.children.entry(c).or_default();
        }
        link.count -= 1;
    }
}

#[test]
fn test() {
    let mut obj = Trie::new();
    obj.insert("apple".to_string());
    obj.insert("apple".to_string());
    assert_eq!(obj.count_words_equal_to("apple".to_string()), 2);
    assert_eq!(obj.count_words_starting_with("app".to_string()), 2);
    obj.erase("apple".to_string());
    assert_eq!(obj.count_words_equal_to("apple".to_string()), 1);
    assert_eq!(obj.count_words_starting_with("app".to_string()), 1);
    obj.erase("apple".to_string());
    assert_eq!(obj.count_words_starting_with("app".to_string()), 0);
}