Exercise: Trie Problems on LeetCode

labuladongOriginalAbout 1274 words

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 apply them directly to all prefix tree-related problems on LeetCode. Let me demonstrate with a few examples.

Optimization Tips

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

For example, the input for LeetCode prefix tree-related problems is limited to lowercase English letters a-z, so the TrieNode does not need to maintain a children array of size 256; setting it to 26 is enough, which can reduce time and space complexity.

Additionally, the Java/cpp code previously provided includes generics, which are not necessary when solving algorithm problems. Removing generics can also offer some efficiency improvements.

208. Implement Trie (Prefix Tree)

Let's 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 problem is from LeetCode 208. Implement Trie (Prefix Tree).

The functions that the problem asks us to implement are actually part of the TrieSet API, so encapsulating a TrieSet can solve this problem directly:

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

The problem is from LeetCode 648. Replace Words.

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

So, you can store the input list of roots dict into a TrieSet, and then directly 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 problem is from LeetCode 211. Design Add and Search Words Data Structure.

The key point of this problem is the search function's wildcard matching, which is essentially the hasKeyWithPattern method we implemented in TrieSet. You can directly apply it 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 examples all use TrieSet. Now, let's look at the problems using TrieMap.

1804. Implement a Trie (Prefix Tree II)

LeetCode problem 1804, "Implement Trie II", can directly use TrieMap. Each inserted word is a key, and the number of insertions is the corresponding value. By reusing the TrieMap API, you can implement the functions required by the 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 involves directly applying templates, it's not very interesting. Let's look at the last problem, which 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.

The problem is from LeetCode 677. Map Sum Pairs.

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 implementation principles and practical problems of the Trie data structure have been covered. If you have reached this point, truly, applause to you. Understanding on paper is not enough; true understanding requires practice. I recommend implementing the above code yourself and going through the problems to gain a deeper understanding of the Trie tree.

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 063. 替换单词 https://leetcode.cn/problems/UhWRSj

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

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

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

剑指 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