Implement Trie II (Prefix Tree)

medium trie design hash map

Problem

Design a prefix tree that supports duplicate words and removal. Implement a Trie class with: insert(word) to add a word (a word may be inserted more than once); countWordsEqualTo(word) to return how many times exactly that word was inserted; countWordsStartingWith(prefix) to return how many inserted words begin with the given prefix; and erase(word) to remove one previous insertion of that word. It is guaranteed that erase is only called on a word that is currently present.

Input

insert "apple", insert "apple", countWordsEqualTo "apple", countWordsStartingWith "app", erase "apple", countWordsEqualTo "apple"

Output2, 2, 1

"apple" was inserted twice, so it counts 2 and the prefix "app" matches 2 words. After erasing one copy, exactly one "apple" remains.

operations (semicolon-separated: insert/equal/prefix/erase word)

class Trie:
    def __init__(self):
        self.root = {"end": 0, "pre": 0, "kids": {}}
    def insert(self, w):
        n = self.root
        for c in w:
            n["kids"].setdefault(c, {"end": 0, "pre": 0, "kids": {}})
            n = n["kids"][c]
            n["pre"] += 1
        n["end"] += 1
    def _walk(self, w):
        n = self.root
        for c in w:
            if c not in n["kids"]:
                return None
            n = n["kids"][c]
        return n
    def countWordsEqualTo(self, w):
        n = self._walk(w)
        return n["end"] if n else 0
    def countWordsStartingWith(self, p):
        n = self._walk(p)
        return n["pre"] if n else 0
    def erase(self, w):
        n = self.root
        for c in w:
            n = n["kids"][c]
            n["pre"] -= 1
        n["end"] -= 1

class Trie {
  constructor() { this.root = { end: 0, pre: 0, kids: {} }; }
  insert(w) {
    let n = this.root;
    for (const c of w) {
      if (!n.kids[c]) n.kids[c] = { end: 0, pre: 0, kids: {} };
      n = n.kids[c];
      n.pre += 1;
    }
    n.end += 1;
  }
  _walk(w) {
    let n = this.root;
    for (const c of w) {
      if (!n.kids[c]) return null;
      n = n.kids[c];
    }
    return n;
  }
  countWordsEqualTo(w) { const n = this._walk(w); return n ? n.end : 0; }
  countWordsStartingWith(p) { const n = this._walk(p); return n ? n.pre : 0; }
  erase(w) {
    let n = this.root;
    for (const c of w) { n = n.kids[c]; n.pre -= 1; }
    n.end -= 1;
  }
}

class Trie {
    private static class Node {
        int end = 0, pre = 0;
        Map<Character, Node> kids = new HashMap<>();
    }
    private Node root = new Node();
    public void insert(String w) {
        Node n = root;
        for (char c : w.toCharArray()) {
            n.kids.computeIfAbsent(c, k -> new Node());
            n = n.kids.get(c);
            n.pre += 1;
        }
        n.end += 1;
    }
    private Node walk(String w) {
        Node n = root;
        for (char c : w.toCharArray()) {
            if (!n.kids.containsKey(c)) return null;
            n = n.kids.get(c);
        }
        return n;
    }
    public int countWordsEqualTo(String w) {
        Node n = walk(w); return n != null ? n.end : 0;
    }
    public int countWordsStartingWith(String p) {
        Node n = walk(p); return n != null ? n.pre : 0;
    }
    public void erase(String w) {
        Node n = root;
        for (char c : w.toCharArray()) { n = n.kids.get(c); n.pre -= 1; }
        n.end -= 1;
    }
}

class Trie {
    struct Node {
        int end = 0, pre = 0;
        unordered_map<char, Node*> kids;
    };
    Node* root = new Node();
    Node* walk(string w) {
        Node* n = root;
        for (char c : w) {
            if (!n->kids.count(c)) return nullptr;
            n = n->kids[c];
        }
        return n;
    }
public:
    void insert(string w) {
        Node* n = root;
        for (char c : w) {
            if (!n->kids.count(c)) n->kids[c] = new Node();
            n = n->kids[c];
            n->pre += 1;
        }
        n->end += 1;
    }
    int countWordsEqualTo(string w) { Node* n = walk(w); return n ? n->end : 0; }
    int countWordsStartingWith(string p) { Node* n = walk(p); return n ? n->pre : 0; }
    void erase(string w) {
        Node* n = root;
        for (char c : w) { n = n->kids[c]; n->pre -= 1; }
        n->end -= 1;
    }
};

Explanation

The original trie stores a single boolean "is this the end of a word" flag. That cannot answer "how many times" questions or support deletion, so here we replace the flag with two integer counters on every node.

The first counter, end, records how many inserted words finish exactly at this node. The second, pre, records how many inserted words pass through this node — that is, how many words have the path-so-far as a prefix. Both are kept in sync as words come and go.

insert walks the word character by character, creating missing child nodes as before, and bumps pre on each node it steps onto. At the final node it bumps end. So a single character of every node along the path is "credited" for this word.

countWordsEqualTo simply walks to the word's last node and returns that node's end; countWordsStartingWith walks to the prefix's last node and returns its pre. If the path does not exist, the answer is zero. erase mirrors insert exactly: it follows the same path and decrements pre on each node and end on the last one, undoing one insertion.

Example: insert apple twice. Every node on the path a→p→p→l→e now has pre = 2, and the final e node has end = 2. countWordsEqualTo("apple") reads that end and returns 2; countWordsStartingWith("app") reads pre at the second p and returns 2. After erase("apple"), every counter on the path drops by one, so countWordsEqualTo("apple") now returns 1.

Time: O(L) per operation Space: O(total characters inserted)