Implement Trie II — Advanced Operations
Solve this Problemwords to insert (duplicates allowed — the same word can be inserted more than once) followed by a list of erasures to remove one occurrence each of, answer one of two questions about a single query: exactly how many currently-present words equal it, or how many currently-present words have it as a prefix?
The trick is giving every trie node two running totals instead of one flag: how many inserted words end exactly at that node, and how many pass through it at all (ending there or continuing further). Inserting a word increments the "passes through" total all along its path and the "ends here" total at the final node; erasing does the same in reverse, but only when a copy genuinely exists to remove — so an erasure can never push either total below zero. Once built, a query is just one walk down to the relevant node, reading off whichever total the question asked about.
Test Case 1:
Test Case 2:
Test Case 3:
Constraints
- ◆
0 ≤ words.length, erasures.length ≤ 1000 - ◆
1 ≤ words[i].length, erasures[i].length ≤ 30 - ◆
0 ≤ query.length ≤ 30 - ◆
words[i], erasures[i], and query consist of lowercase English letters - ◆
an erasure applied when nothing matching remains to erase is simply a no-op
Try the Dry Run
Don't just read the solution — watch it execute, one step at a time.
Approach & Solutions
Brute Force — Hash Map of Word Counts
BruteKeep a hash map from word to how many times it's currently present. Inserting bumps a count; erasing decrements one (only if a copy is actually there, so an erasure never drives a count negative). Answering countWordsEqualTo is then one O(1) lookup — but countWordsStartingWith has to walk every distinct word (D of them, each up to length L) checking whether it starts with the query, since nothing about the map itself groups words by shared prefixes.
O(D · L)O(D · L)1class Solution {
2 public int trieAdvancedQuery(String[] words, String[] erasures, String query, boolean isPrefixQuery) {
3 Map<String, Integer> counts = new HashMap<>();
4 for (String word : words) {
5 counts.merge(word, 1, Integer::sum);
6 }
7 for (String word : erasures) {
8 Integer c = counts.get(word);
9 if (c != null && c > 0) {
10 counts.put(word, c - 1);
11 }
12 }
13 int total = 0;
14 for (Map.Entry<String, Integer> entry : counts.entrySet()) {
15 if (isPrefixQuery) {
16 if (entry.getKey().startsWith(query)) total += entry.getValue();
17 } else {
18 if (entry.getKey().equals(query)) total += entry.getValue();
19 }
20 }
21 return total;
22 }
23}Optimal — Trie with Counters
OptimalGive every trie node two counters instead of a single isEnd flag: wordCount (how many inserted words end exactly here) and prefixCount (how many inserted words pass through here at all, end or not). Inserting a word bumps prefixCount along its whole path and wordCount at the last node; erasing does the reverse — but only after confirming a copy of that word genuinely exists, so counts never go negative. Once built, both kinds of query are just one O(L) walk down to the query's node, reading off whichever counter the query asked for.
O(L)O(N)1class Solution {
2 static class TrieNode {
3 TrieNode[] children = new TrieNode[26];
4 int wordCount = 0;
5 int prefixCount = 0;
6 }
7
8 private void insert(TrieNode root, String word) {
9 TrieNode node = root;
10 node.prefixCount++;
11 for (char c : word.toCharArray()) {
12 int idx = c - 'a';
13 if (node.children[idx] == null) {
14 node.children[idx] = new TrieNode();
15 }
16 node = node.children[idx];
17 node.prefixCount++;
18 }
19 node.wordCount++;
20 }
21
22 private boolean canErase(TrieNode root, String word) {
23 TrieNode node = root;
24 for (char c : word.toCharArray()) {
25 int idx = c - 'a';
26 if (node.children[idx] == null) return false;
27 node = node.children[idx];
28 }
29 return node.wordCount > 0;
30 }
31
32 private void erase(TrieNode root, String word) {
33 if (!canErase(root, word)) return;
34 TrieNode node = root;
35 node.prefixCount--;
36 for (char c : word.toCharArray()) {
37 int idx = c - 'a';
38 node = node.children[idx];
39 node.prefixCount--;
40 }
41 node.wordCount--;
42 }
43
44 public int trieAdvancedQuery(String[] words, String[] erasures, String query, boolean isPrefixQuery) {
45 TrieNode root = new TrieNode();
46 for (String word : words) insert(root, word);
47 for (String word : erasures) erase(root, word);
48
49 TrieNode node = root;
50 for (char c : query.toCharArray()) {
51 int idx = c - 'a';
52 if (node.children[idx] == null) return 0;
53 node = node.children[idx];
54 }
55 return isPrefixQuery ? node.prefixCount : node.wordCount;
56 }
57}