Implement Trie ll

Medium Tries Basics Original on Code360

Implement Trie ll is a medium tries problem solved with the basics pattern. The best approach, optimal (trie with prefix and end counters), runs in O(L) per operation time and O(total characters · 26) space. Below are 2 approaches in Java, from hash map of word counts up.

Problem

Implement a trie that also counts: insert(word), countWordsEqualTo(word), countWordsStartingWith(prefix) and erase(word) (remove one copy). Duplicate words are allowed.

Examples

Example 1

Input: insert(apple), insert(apple), insert(apps), countWordsEqualTo(apple), countWordsStartingWith(app), erase(apple), countWordsEqualTo(apple)
Output: 2, 3, 1

Constraints

Lowercase words; erase is only called on words that exist.
Up to 3 * 10^4 operations.

Animated walkthrough

A narrated, step-by-step animation that builds the solution from the idea up. Press play, or step through it at your own pace.

Concepts first, then the problem and every approach, step by step.

Code

Code appears here once we start writing it.

Space to play or pause · ← → to jump a step · click or drag the bar to seek

Solutions

Try it yourself first. Then compare: each approach lists its idea, the steps, its time and space, and the Java code.

Approach	Time	Space
Hash map of word counts	`O(n · L) for prefix counts`	`O(total characters)`
Optimal (trie with prefix and end counters)	`O(L) per operation`	`O(total characters · 26)`

1Hash map of word counts

TimeO(n · L) for prefix counts

SpaceO(total characters)

Keep a map word → count. Counting prefixes scans every word.

insert/erase adjust the count.
countWordsStartingWith sums counts of words with the prefix.

Java

class Trie {
    private final Map<String, Integer> count = new HashMap<>();

    public void insert(String word) { count.merge(word, 1, Integer::sum); }

    public int countWordsEqualTo(String word) { return count.getOrDefault(word, 0); }

    public int countWordsStartingWith(String prefix) {
        int total = 0;
        for (Map.Entry<String, Integer> e : count.entrySet())
            if (e.getKey().startsWith(prefix)) total += e.getValue();
        return total;
    }

    public void erase(String word) {
        if (count.merge(word, -1, Integer::sum) == 0) count.remove(word);
    }
}

2Optimal (trie with prefix and end counters)

TimeO(L) per operation

SpaceO(total characters · 26)

Each node stores prefix, the number of words passing through it, and end, the number of words ending there. insert increments along the path, erase decrements along it, and both counts are read at the last node of the walk.

insert: for each char, create if missing, move, prefix++. Then end++.
countWordsEqualTo: walk; return end or 0.
countWordsStartingWith: walk; return prefix or 0.
erase: walk, prefix-- at each node; end-- at the last.

Java

class Trie {
    private static class Node {
        Node[] next = new Node[26];
        int prefix, end;
    }

    private final Node root = new Node();

    public void insert(String word) {
        Node cur = root;
        for (char c : word.toCharArray()) {
            if (cur.next[c - 'a'] == null) cur.next[c - 'a'] = new Node();
            cur = cur.next[c - 'a'];
            cur.prefix++;
        }
        cur.end++;
    }

    public int countWordsEqualTo(String word) {
        Node n = walk(word);
        return n == null ? 0 : n.end;
    }

    public int countWordsStartingWith(String prefix) {
        Node n = walk(prefix);
        return n == null ? 0 : n.prefix;
    }

    public void erase(String word) {
        Node cur = root;
        for (char c : word.toCharArray()) {
            cur = cur.next[c - 'a'];
            cur.prefix--;
        }
        cur.end--;
    }

    private Node walk(String s) {
        Node cur = root;
        for (char c : s.toCharArray()) {
            cur = cur.next[c - 'a'];
            if (cur == null) return null;
        }
        return cur;
    }
}

Edge cases to test

Inserting the same word several times
Erasing one copy of a duplicated word

Hints

Hint 1

Store two counters in every node: how many words pass through it (prefix count) and how many words end at it.

FAQ

What is the best time complexity for Implement Trie ll?

Optimal (trie with prefix and end counters) runs in O(L) per operation time and O(total characters · 26) extra space.

Which pattern does Implement Trie ll use?

It is a tries problem that uses the basics pattern. Other problems with the same pattern: Implement Trie (Prefix Tree).

Is there a brute force solution for Implement Trie ll?

Yes. Hash map of word counts takes O(n · L) for prefix counts time and O(total characters) space. Keep a map word → count.

Which edge cases should I test for Implement Trie ll?

Inserting the same word several times; Erasing one copy of a duplicated word.