Exercise: Trie Problems on LeetCode

labuladongOriginalAbout 1322 words

Note

Now all the plugins has supported English. I'm still improving the website...

This article will resolve

LeetCode	Difficulty
1804. Implement Trie II (Prefix Tree)🔒	🟠
208. Implement Trie (Prefix Tree)	🟠
211. Design Add and Search Words Data Structure	🟠
648. Replace Words	🟠
677. Map Sum Pairs	🟠

前置知识

阅读本文前，你需要先学习：

With TrieMap and TrieSet, you can directly apply them to all prefix tree-related problems on LeetCode. Let me demonstrate with a few examples.

Optimization Tips

Firstly, the TrieMap/TrieSet implementation provided in the previous article TrieMap/TrieSet Code Implementation has room for optimization in specific problems.

For instance, the input for prefix tree-related problems on LeetCode is restricted to lowercase English letters a-z. Therefore, the TrieNode does not need to maintain a children array of size 256; a size of 26 is sufficient, reducing time and space complexity.

Additionally, the Java/Cpp code provided earlier includes generics, which are not necessary when solving algorithm problems. Removing generics can also improve efficiency.

208. Implement Trie (Prefix Tree)

Let's take a look at LeetCode problem 208, "Implement Trie (Prefix Tree)":

208. Implement Trie (Prefix Tree) | LeetCode |

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.
boolean search(String word) Returns true if the string word is in the trie (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.

Example 1:

Input
["Trie", "insert", "search", "search", "startsWith", "insert", "search"]
[[], ["apple"], ["apple"], ["app"], ["app"], ["app"], ["app"]]
Output
[null, null, true, false, true, null, true]

Explanation
Trie trie = new Trie();
trie.insert("apple");
trie.search("apple");   // return True
trie.search("app");     // return False
trie.startsWith("app"); // return True
trie.insert("app");
trie.search("app");     // return True

Constraints:

1 <= word.length, prefix.length <= 2000
word and prefix consist only of lowercase English letters.
At most 3 * 10⁴ calls in total will be made to insert, search, and startsWith.

The functions that the problem asks us to implement are actually part of the TrieSet API. So, we can solve this problem by simply encapsulating a TrieSet:

class Trie {
    // directly encapsulate TrieSet
    TrieSet set = new TrieSet();
    
    // insert an element
    public void insert(String word) {
        set.add(word);
    }
    
    // check if the element is in the set
    public boolean search(String word) {
        return set.contains(word);
    }
    
    // check if there is an element with prefix in the set
    public boolean startsWith(String prefix) {
        return set.hasKeyWithPrefix(prefix);
    }
}

// see above
class TrieSet {}

648. Replace Words

Next, let's look at LeetCode's Problem 648 "Replace Words":

648. Replace Words | LeetCode |

In English, we have a concept called root, which can be followed by some other word to form another longer word - let's call this word derivative. For example, when the root "help" is followed by the word "ful", we can form a derivative "helpful".

Given a dictionary consisting of many roots and a sentence consisting of words separated by spaces, replace all the derivatives in the sentence with the root forming it. If a derivative can be replaced by more than one root, replace it with the root that has the shortest length.

Return the sentence after the replacement.

Example 1:

Input: dictionary = ["cat","bat","rat"], sentence = "the cattle was rattled by the battery"
Output: "the cat was rat by the bat"

Example 2:

Input: dictionary = ["a","b","c"], sentence = "aadsfasf absbs bbab cadsfafs"
Output: "a a b c"

Constraints:

1 <= dictionary.length <= 1000
1 <= dictionary[i].length <= 100
dictionary[i] consists of only lower-case letters.
1 <= sentence.length <= 10⁶
sentence consists of only lower-case letters and spaces.
The number of words in sentence is in the range [1, 1000]
The length of each word in sentence is in the range [1, 1000]
Every two consecutive words in sentence will be separated by exactly one space.
sentence does not have leading or trailing spaces.

Now that you have learned the Trie data structure, you should be able to see that this problem is about finding the shortest prefix.

So you can store the input list of roots dict in a TrieSet, and then simply reuse the shortestPrefixOf function we implemented:

String replaceWords(List<String> dict, String sentence) {
    // first, store all the roots in the TrieSet
    TrieSet set = new TrieSet();
    for (String key : dict) {
        set.add(key);
    }
    StringBuilder sb = new StringBuilder();
    String[] words = sentence.split(" ");
    // process the words in the sentence
    for (int i = 0; i < words.length; i++) {
        // search for the shortest root (shortest prefix) in the Trie tree
        String prefix = set.shortestPrefixOf(words[i]);
        if (!prefix.isEmpty()) {
            // if found, replace with the root
            sb.append(prefix);
        } else {
            // otherwise, keep the original word
            sb.append(words[i]);
        }
        
        if (i != words.length - 1) {
            // add spaces between words
            sb.append(' ');
        }
    }

    return sb.toString();
}

// see above
class TrieSet {}

// see above
class TrieMap {}

211. Add and Search Word - Data Structure Design

Let's continue with LeetCode Problem 211: "Add and Search Word - Data Structure Design":

211. Design Add and Search Words Data Structure | LeetCode |

Design a data structure that supports adding new words and finding if a string matches any previously added string.

Implement the WordDictionary class:

WordDictionary() Initializes the object.
void addWord(word) Adds word to the data structure, it can be matched later.
bool search(word) Returns true if there is any string in the data structure that matches word or false otherwise. word may contain dots '.' where dots can be matched with any letter.

Example:

Input
["WordDictionary","addWord","addWord","addWord","search","search","search","search"]
[[],["bad"],["dad"],["mad"],["pad"],["bad"],[".ad"],["b.."]]
Output
[null,null,null,null,false,true,true,true]

Explanation
WordDictionary wordDictionary = new WordDictionary();
wordDictionary.addWord("bad");
wordDictionary.addWord("dad");
wordDictionary.addWord("mad");
wordDictionary.search("pad"); // return False
wordDictionary.search("bad"); // return True
wordDictionary.search(".ad"); // return True
wordDictionary.search("b.."); // return True

Constraints:

1 <= word.length <= 25
word in addWord consists of lowercase English letters.
word in search consist of '.' or lowercase English letters.
There will be at most 2 dots in word for search queries.
At most 10⁴ calls will be made to addWord and search.

The key point of this problem lies in the search function for wildcard matching. In fact, it's similar to the hasKeyWithPattern method we implemented in TrieSet. You can directly apply that here:

class WordDictionary {
    TrieSet set = new TrieSet();
    
    // add elements to TrieSet
    public void addWord(String word) {
        set.add(word);
    }
    
    // wildcard matching elements
    public boolean search(String word) {
        return set.hasKeyWithPattern(word);
    }
}

// see above
class TrieSet {}

// see above
class TrieMap {}

The above problems all use TrieSet. Now let's look at problems involving TrieMap.

1804. Implement a Trie (Prefix Tree II)

First, let's look at LeetCode problem 1804 "Implement Trie II":

1804. Implement Trie II (Prefix Tree) | LeetCode |

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⁴ 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.

This problem can use TrieMap, where each inserted word is the key, and the number of insertions is the corresponding value. By reusing the TrieMap API, you can implement the required functions for this problem:

class Trie {
    // encapsulate our implementation of TrieMap
    TrieMap<Integer> map = new TrieMap<>();
    
    // insert word and record the number of insertions
    public void insert(String word) {
        if (!map.containsKey(word)) {
            map.put(word, 1);
        } else {
            map.put(word, map.get(word) + 1);
        }
    }
    
    // query the number of insertions of word
    public int countWordsEqualTo(String word) {
        if (!map.containsKey(word)) {
            return 0;
        }
        return map.get(word);
    }
    
    // accumulate the total number of insertions for keys with prefix
    public int countWordsStartingWith(String prefix) {
        int res = 0;
        for (String key : map.keysWithPrefix(prefix)) {
            res += map.get(key);
        }
        return res;
    }
    
    // decrease the number of insertions of word by one
    public void erase(String word) {
        int freq = map.get(word);
        if (freq - 1 == 0) {
            map.remove(word);
        } else {
            map.put(word, freq - 1);
        }
    }
}

// see above
class TrieMap {}

677. Map Sum Pairs

Since it's just applying a template directly, let's look at the last problem. This is LeetCode Problem 677: "Map Sum Pairs":

677. Map Sum Pairs | LeetCode |

Design a map that allows you to do the following:

Maps a string key to a given value.
Returns the sum of the values that have a key with a prefix equal to a given string.

Implement the MapSum class:

MapSum() Initializes the MapSum object.
void insert(String key, int val) Inserts the key-val pair into the map. If the key already existed, the original key-value pair will be overridden to the new one.
int sum(string prefix) Returns the sum of all the pairs' value whose key starts with the prefix.

Example 1:

Input
["MapSum", "insert", "sum", "insert", "sum"]
[[], ["apple", 3], ["ap"], ["app", 2], ["ap"]]
Output
[null, null, 3, null, 5]

Explanation
MapSum mapSum = new MapSum();
mapSum.insert("apple", 3);  
mapSum.sum("ap");           // return 3 (apple = 3)
mapSum.insert("app", 2);    
mapSum.sum("ap");           // return 5 (apple + app = 3 + 2 = 5)

Constraints:

1 <= key.length, prefix.length <= 50
key and prefix consist of only lowercase English letters.
1 <= val <= 1000
At most 50 calls will be made to insert and sum.

This problem is a standard application of TrieMap. Let's directly look at the code:

class MapSum {
    // encapsulate our implemented TrieMap
    TrieMap<Integer> map = new TrieMap<>();
    
    // insert key-value pair
    public void insert(String key, int val) {
        map.put(key, val);
    }
    
    // sum the values of all keys with the prefix
    public int sum(String prefix) {
        List<String> keys = map.keysWithPrefix(prefix);
        int res = 0;
        for (String key : keys) {
            res += map.get(key);
        }
        return res;
    }
}

// see above
class TrieMap {}

The principles and practical problems of implementing the Trie data structure have been covered. If you've read this far, you truly deserve applause. There's a saying, "Learning from books is superficial; true understanding comes from practice." I highly recommend that you personally implement the code above and solve related problems to gain a deeper understanding of the Trie tree.

Related Problems

You can install my Chrome extension then open the link.

LeetCode	Difficulty
14. Longest Common Prefix	🟢

Exercise: Trie Problems on LeetCode

1804. 实现 Trie （前缀树） II https://leetcode.cn/problems/implement-trie-ii-prefix-tree

1804. Implement Trie II (Prefix Tree) https://leetcode.com/problems/implement-trie-ii-prefix-tree

208. 实现 Trie (前缀树) https://leetcode.cn/problems/implement-trie-prefix-tree

208. Implement Trie (Prefix Tree) https://leetcode.com/problems/implement-trie-prefix-tree

211. 添加与搜索单词 - 数据结构设计 https://leetcode.cn/problems/design-add-and-search-words-data-structure

211. Design Add and Search Words Data Structure https://leetcode.com/problems/design-add-and-search-words-data-structure

648. 单词替换 https://leetcode.cn/problems/replace-words

648. Replace Words https://leetcode.com/problems/replace-words

677. 键值映射 https://leetcode.cn/problems/map-sum-pairs

677. Map Sum Pairs https://leetcode.com/problems/map-sum-pairs

剑指 Offer II 062. 实现前缀树 https://leetcode.cn/problems/QC3q1f

剑指 Offer II 062. 实现前缀树 https://leetcode.com/problems/QC3q1f

剑指 Offer II 063. 替换单词 https://leetcode.cn/problems/UhWRSj

剑指 Offer II 063. 替换单词 https://leetcode.com/problems/UhWRSj

剑指 Offer II 066. 单词之和 https://leetcode.cn/problems/z1R5dt

剑指 Offer II 066. 单词之和 https://leetcode.com/problems/z1R5dt

208. Implement Trie (Prefix Tree)

648. Replace Words

211. Add and Search Word - Data Structure Design

1804. Implement a Trie (Prefix Tree II)

677. Map Sum Pairs