☕ DSA Java Articles

Topics › Tries › Problems

Trie Implementation and Advanced Operations

Hard TUF ID 1027 A2ZSDEBlind 75

LeetCode ↗ Java source ↗ Edit this article ↗

Problem

Design a structure that stores a multiset of words (the same word may be inserted several times) and supports four operations:

This is “Implement Trie II” (LeetCode 1804). Constraints there: lowercase letters, lengths from 1 to 2000, at most 3 × 10^4 calls, and erase is only called for a word that is present. The code below is a little more forgiving: erasing a word that is not stored does nothing, and the empty string works like any other word.

Example 1

insert("apple"), insert("apple")
countWordsEqualTo("apple")       -> 2
countWordsStartingWith("app")    -> 2
erase("apple")
countWordsEqualTo("apple")       -> 1
countWordsStartingWith("app")    -> 1
erase("apple")
countWordsStartingWith("app")    -> 0     // nothing left

Example 2

insert("ab"), insert("abc"), insert("abcd")
countWordsStartingWith("ab")     -> 3     // "ab" itself counts
countWordsStartingWith("abc")    -> 2
erase("abc")
countWordsStartingWith("ab")     -> 2
countWordsEqualTo("abcd")        -> 1     // erasing "abc" must not disturb "abcd"

All three designs implement the same interface, so the Java file drives them with identical scripts and compares every answer:

interface WordCounter {
    void insert(String word);
    int countWordsEqualTo(String word);
    int countWordsStartingWith(String prefix);
    void erase(String word);
}

Approach 1: Brute force, a list of words

Intuition. Keep every inserted copy in a list. Both counting questions become a scan with a counter: count the words that are equal to the query, or that startsWith it. Erasing one copy is List.remove(Object), which deletes the first equal element and quietly returns false if there is none. Correct and short, but every query touches the whole collection.

Algorithm.

  1. insert: append the word.
  2. countWordsEqualTo: scan and count w.equals(word).
  3. countWordsStartingWith: scan and count w.startsWith(prefix).
  4. erase: remove the first equal element, if any.
static class ListCounter implements WordCounter {
    private final List<String> words = new ArrayList<>();

    public void insert(String word) {
        words.add(word);
    }

    public int countWordsEqualTo(String word) {
        int count = 0;
        for (String w : words) {
            if (w.equals(word)) count++;
        }
        return count;
    }

    public int countWordsStartingWith(String prefix) {
        int count = 0;
        for (String w : words) {
            if (w.startsWith(prefix)) count++;
        }
        return count;
    }

    public void erase(String word) {
        words.remove(word);                            // removes one copy; does nothing if the word is absent
    }
}

Complexity. insert is O(1) amortised. Each count and each erase is O(n · L) for n stored words, since every comparison may read the whole query. Space is O(total characters).

Approach 2: Better, hash maps of word counts and prefix counts

Intuition. Instead of recomputing the answers on every query, maintain them. Keep wordCount[w], the number of copies of each word, and prefixCount[p], the number of stored words that begin with p. When a word arrives, every one of its prefixes gains one word, so we bump L + 1 entries: the prefixes of lengths 0 through L. Including the empty prefix makes countWordsStartingWith("") equal to the total. Erasing does the same with -1, but only if the word is actually present; otherwise we would corrupt counts of words that merely share a prefix with it.

computeIfPresent(key, (k, v) -> v == 1 ? null : v - 1) decrements and removes an entry that would drop to zero, so the maps never fill up with stale zeros.

Algorithm.

  1. insert(word): wordCount[word]++; for len = 0..L, prefixCount[word[0..len)]++.
  2. Queries: look the string up in the matching map, defaulting to 0.
  3. erase(word): if wordCount[word] is 0, stop. Otherwise decrement wordCount[word] and all L + 1 prefix entries.
static class MapCounter implements WordCounter {
    private final Map<String, Integer> wordCount = new HashMap<>();
    private final Map<String, Integer> prefixCount = new HashMap<>();

    public void insert(String word) {
        wordCount.merge(word, 1, Integer::sum);
        for (int len = 0; len <= word.length(); len++) {
            prefixCount.merge(word.substring(0, len), 1, Integer::sum);
        }
    }

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

    public int countWordsStartingWith(String prefix) {
        return prefixCount.getOrDefault(prefix, 0);
    }

    public void erase(String word) {
        if (countWordsEqualTo(word) == 0) return;      // nothing to erase: leave every count untouched
        wordCount.computeIfPresent(word, (k, v) -> v == 1 ? null : v - 1);
        for (int len = 0; len <= word.length(); len++) {
            prefixCount.computeIfPresent(word.substring(0, len), (k, v) -> v == 1 ? null : v - 1);
        }
    }
}

Complexity. insert and erase are O(L^2): L + 1 substrings of average length L / 2 are created and hashed. Both queries are O(L) on average. Space is O(L^2) characters per distinct word, because every prefix is stored as its own string.

The prefixes of one word overlap almost completely, and prefixes are shared between words, yet each map entry stores its own copy. A trie gives each distinct prefix exactly one node, which is the natural home for its counter.

Approach 3: Optimal, a trie with two counters per node

Intuition. In a trie the node reached by walking p from the root is the prefix p. So the counter prefixCount[p] from Approach 2 can live on that node, and walking a word visits exactly its L + 1 prefix nodes, in O(1) each, with no strings built. The basic trie’s boolean isEnd becomes a counter too, because duplicates must be counted. Each node keeps:

The root’s prefixCount is the total number of stored words, which is exactly the answer for the empty prefix.

Algorithm.

  1. insert(word): increment root.prefixCount. For each letter, create the child if missing, move to it, and increment its prefixCount. At the last node increment endCount.
  2. countWordsEqualTo(word): walk; a missing edge means 0, otherwise return endCount.
  3. countWordsStartingWith(prefix): walk; a missing edge means 0, otherwise return prefixCount.
  4. erase(word): if countWordsEqualTo(word) is 0, stop. Otherwise decrement root.prefixCount, then walk the word decrementing each node’s prefixCount. If a node’s count reaches 0, unlink it from its parent and stop; if the walk completes, decrement endCount on the last node.
static class TrieCounter implements WordCounter {
    private static class Node {
        final Node[] child = new Node[26];
        int endCount;                                  // words that end exactly at this node
        int prefixCount;                               // words that pass through or end at this node
    }

    private final Node root = new Node();

    public void insert(String word) {
        Node cur = root;
        cur.prefixCount++;                             // every word has the empty prefix
        for (int i = 0; i < word.length(); i++) {
            int c = word.charAt(i) - 'a';
            if (cur.child[c] == null) cur.child[c] = new Node();
            cur = cur.child[c];
            cur.prefixCount++;
        }
        cur.endCount++;
    }

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

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

    public void erase(String word) {
        if (countWordsEqualTo(word) == 0) return;      // nothing to erase: leave every count untouched
        Node cur = root;
        cur.prefixCount--;
        for (int i = 0; i < word.length(); i++) {
            int c = word.charAt(i) - 'a';
            Node next = cur.child[c];
            if (--next.prefixCount == 0) {             // no word uses this branch any more: drop all of it
                cur.child[c] = null;
                return;
            }
            cur = next;
        }
        cur.endCount--;
    }

    /** Follows s from the root; returns the node reached, or null as soon as a letter has no edge. */
    private Node walk(String s) {
        Node cur = root;
        for (int i = 0; i < s.length(); i++) {
            cur = cur.child[s.charAt(i) - 'a'];
            if (cur == null) return null;
        }
        return cur;
    }
}

Why unlinking is safe. A node’s prefixCount is the number of words in its whole subtree. Once it is 0, every counter below it is 0 as well, so dropping the subtree changes no answer; it only frees memory, keeping the trie proportional to the words currently stored rather than every word ever inserted. The walk must return right after unlinking, because the next node no longer exists.

Why the presence check matters. Without it, erase("app") when only "apple" is stored would decrement the counters on a, p, p and set the endCount of "app" to -1. countWordsStartingWith("app") would then report 0 even though "apple" is still there.

Dry run on Example 2, showing prefixCount on each node (blank means the node does not exist yet):

Step root a ab abc abcd endCount change
insert("ab") 1 1 1     ab: 1
insert("abc") 2 2 2 1   abc: 1
insert("abcd") 3 3 3 2 1 abcd: 1
erase("abc") 2 2 2 1 1 abc: 0

Before the erase, countWordsStartingWith("ab") reads 3 and countWordsStartingWith("abc") reads 2. After it, "ab" reads 2, and countWordsEqualTo("abcd") still reads the untouched 1.

Complexity. Every operation is O(L) time: one walk, constant work per node. Space is O(N · 26) references for N stored characters, with shared prefixes stored once.

Edge cases and pitfalls