Implement Trie II (Prefix Tree)

Problem

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.
• int countWordsEqualTo(String word) Returns the number of instances of the string word in the trie.
• int countWordsStartingWith(String prefix) Returns the number of strings in the trie that have the string prefix as a prefix.
• void erase(String word) Erases the string word from the trie.

Examples

Example 1:

**Input**
["Trie", "insert", "insert", "countWordsEqualTo", "countWordsStartingWith", "erase", "countWordsEqualTo", "countWordsStartingWith", "erase", "countWordsStartingWith"]
[[], ["apple"], ["apple"], ["apple"], ["app"], ["apple"], ["apple"], ["app"], ["apple"], ["app"]]
**Output**
[null, null, null, 2, 2, null, 1, 1, null, 0]
**Explanation**
Trie trie = new Trie();
trie.insert("apple");               // Inserts "apple".
trie.insert("apple");               // Inserts another "apple".
trie.countWordsEqualTo("apple");    // There are two instances of "apple" so return 2.
trie.countWordsStartingWith("app"); // "app" is a prefix of "apple" so return 2.
trie.erase("apple");                // Erases one "apple".
trie.countWordsEqualTo("apple");    // Now there is only one instance of "apple" so return 1.
trie.countWordsStartingWith("app"); // return 1
trie.erase("apple");                // Erases "apple". Now the trie is empty.
trie.countWordsStartingWith("app"); // return 0

Constraints:

• 1 <= word.length, prefix.length <= 2000
• word and prefix consist only of lowercase English letters.
• At most 3 * 10^4 calls in total will be made to insert, countWordsEqualTo, countWordsStartingWith, and erase.
• It is guaranteed that for any function call to erase, the string word will exist in the trie.

Solution

Method 1 – Trie with Count Tracking

Intuition

We need to support insert, erase, and two types of count queries efficiently. A Trie (prefix tree) is ideal for prefix-based operations, and we can augment each node to track the number of words passing through and ending at that node.

Approach

1. Each Trie node stores:
   • children: mapping for next characters.
   • cnt: number of words passing through this node (for prefix count).
   • end: number of words ending at this node (for exact word count).
2. For insert, traverse the Trie, incrementing cnt at each node and end at the last node.
3. For countWordsEqualTo, traverse to the end node and return end.
4. For countWordsStartingWith, traverse to the prefix node and return cnt.
5. For erase, traverse the Trie, decrementing cnt at each node and end at the last node.

Code

C++

class Trie {
    struct Node {
        Node* nxt[26] = {};
        int cnt = 0, end = 0;
    } *root;
public:
    Trie() { root = new Node(); }
    void insert(string w) {
        Node* cur = root;
        for (char c : w) {
            int i = c - 'a';
            if (!cur->nxt[i]) cur->nxt[i] = new Node();
            cur = cur->nxt[i];
            cur->cnt++;
        }
        cur->end++;
    }
    int countWordsEqualTo(string w) {
        Node* cur = root;
        for (char c : w) {
            int i = c - 'a';
            if (!cur->nxt[i]) return 0;
            cur = cur->nxt[i];
        }
        return cur->end;
    }
    int countWordsStartingWith(string p) {
        Node* cur = root;
        for (char c : p) {
            int i = c - 'a';
            if (!cur->nxt[i]) return 0;
            cur = cur->nxt[i];
        }
        return cur->cnt;
    }
    void erase(string w) {
        Node* cur = root;
        for (char c : w) {
            int i = c - 'a';
            cur = cur->nxt[i];
            cur->cnt--;
        }
        cur->end--;
    }
};

Go

type Node struct {
    nxt [26]*Node
    cnt, end int
}

type Trie struct {
    root *Node
}

func Constructor() Trie {
    return Trie{root: &Node{}}
}

func (t *Trie) Insert(w string) {
    cur := t.root
    for _, c := range w {
        i := c - 'a'
        if cur.nxt[i] == nil {
            cur.nxt[i] = &Node{}
        }
        cur = cur.nxt[i]
        cur.cnt++
    }
    cur.end++
}

func (t *Trie) CountWordsEqualTo(w string) int {
    cur := t.root
    for _, c := range w {
        i := c - 'a'
        if cur.nxt[i] == nil {
            return 0
        }
        cur = cur.nxt[i]
    }
    return cur.end
}

func (t *Trie) CountWordsStartingWith(p string) int {
    cur := t.root
    for _, c := range p {
        i := c - 'a'
        if cur.nxt[i] == nil {
            return 0
        }
        cur = cur.nxt[i]
    }
    return cur.cnt
}

func (t *Trie) Erase(w string) {
    cur := t.root
    for _, c := range w {
        i := c - 'a'
        cur = cur.nxt[i]
        cur.cnt--
    }
    cur.end--
}

Java

class Trie {
    class Node {
        Node[] nxt = new Node[26];
        int cnt = 0, end = 0;
    }
    Node root = new Node();
    public void insert(String w) {
        Node cur = root;
        for (char c : w.toCharArray()) {
            int i = c - 'a';
            if (cur.nxt[i] == null) cur.nxt[i] = new Node();
            cur = cur.nxt[i];
            cur.cnt++;
        }
        cur.end++;
    }
    public int countWordsEqualTo(String w) {
        Node cur = root;
        for (char c : w.toCharArray()) {
            int i = c - 'a';
            if (cur.nxt[i] == null) return 0;
            cur = cur.nxt[i];
        }
        return cur.end;
    }
    public int countWordsStartingWith(String p) {
        Node cur = root;
        for (char c : p.toCharArray()) {
            int i = c - 'a';
            if (cur.nxt[i] == null) return 0;
            cur = cur.nxt[i];
        }
        return cur.cnt;
    }
    public void erase(String w) {
        Node cur = root;
        for (char c : w.toCharArray()) {
            int i = c - 'a';
            cur = cur.nxt[i];
            cur.cnt--;
        }
        cur.end--;
    }
}

Kotlin

class Trie {
    class Node {
        val nxt = Array<Node?>(26) { null }
        var cnt = 0
        var end = 0
    }
    private val root = Node()
    fun insert(w: String) {
        var cur = root
        for (c in w) {
            val i = c - 'a'
            if (cur.nxt[i] == null) cur.nxt[i] = Node()
            cur = cur.nxt[i]!!
            cur.cnt++
        }
        cur.end++
    }
    fun countWordsEqualTo(w: String): Int {
        var cur = root
        for (c in w) {
            val i = c - 'a'
            if (cur.nxt[i] == null) return 0
            cur = cur.nxt[i]!!
        }
        return cur.end
    }
    fun countWordsStartingWith(p: String): Int {
        var cur = root
        for (c in p) {
            val i = c - 'a'
            if (cur.nxt[i] == null) return 0
            cur = cur.nxt[i]!!
        }
        return cur.cnt
    }
    fun erase(w: String) {
        var cur = root
        for (c in w) {
            val i = c - 'a'
            cur = cur.nxt[i]!!
            cur.cnt--
        }
        cur.end--
    }
}

Python

class TrieNode:
    def __init__(self):
        self.nxt = [None] * 26
        self.cnt = 0
        self.end = 0

class Solution:
    def __init__(self):
        self.root = TrieNode()
    def insert(self, w: str) -> None:
        cur = self.root
        for c in w:
            i = ord(c) - ord('a')
            if not cur.nxt[i]:
                cur.nxt[i] = TrieNode()
            cur = cur.nxt[i]
            cur.cnt += 1
        cur.end += 1
    def countWordsEqualTo(self, w: str) -> int:
        cur = self.root
        for c in w:
            i = ord(c) - ord('a')
            if not cur.nxt[i]:
                return 0
            cur = cur.nxt[i]
        return cur.end
    def countWordsStartingWith(self, p: str) -> int:
        cur = self.root
        for c in p:
            i = ord(c) - ord('a')
            if not cur.nxt[i]:
                return 0
            cur = cur.nxt[i]
        return cur.cnt
    def erase(self, w: str) -> None:
        cur = self.root
        for c in w:
            i = ord(c) - ord('a')
            cur = cur.nxt[i]
            cur.cnt -= 1
        cur.end -= 1

Rust

struct Node {
    nxt: [Option<Box<Node>>; 26],
    cnt: i32,
    end: i32,
}

impl Node {
    fn new() -> Self {
        Self { nxt: Default::default(), cnt: 0, end: 0 }
    }
}

struct Solution {
    root: Box<Node>,
}

impl Solution {
    fn new() -> Self {
        Self { root: Box::new(Node::new()) }
    }
    fn insert(&mut self, w: &str) {
        let mut cur = &mut self.root;
        for c in w.chars() {
            let i = (c as u8 - b'a') as usize;
            cur = cur.nxt[i].get_or_insert_with(|| Box::new(Node::new()));
            cur.cnt += 1;
        }
        cur.end += 1;
    }
    fn count_words_equal_to(&self, w: &str) -> i32 {
        let mut cur = &self.root;
        for c in w.chars() {
            let i = (c as u8 - b'a') as usize;
            if let Some(ref n) = cur.nxt[i] {
                cur = n;
            } else {
                return 0;
            }
        }
        cur.end
    }
    fn count_words_starting_with(&self, p: &str) -> i32 {
        let mut cur = &self.root;
        for c in p.chars() {
            let i = (c as u8 - b'a') as usize;
            if let Some(ref n) = cur.nxt[i] {
                cur = n;
            } else {
                return 0;
            }
        }
        cur.cnt
    }
    fn erase(&mut self, w: &str) {
        let mut cur = &mut self.root;
        for c in w.chars() {
            let i = (c as u8 - b'a') as usize;
            cur = cur.nxt[i].as_deref_mut().unwrap();
            cur.cnt -= 1;
        }
        cur.end -= 1;
    }
}

TypeScript

class TrieNode {
  nxt: (TrieNode | null)[] = Array(26).fill(null);
  cnt = 0;
  end = 0;
}

class Solution {
  root = new TrieNode();
  insert(w: string): void {
    let cur = this.root;
    for (const c of w) {
      const i = c.charCodeAt(0) - 97;
      if (!cur.nxt[i]) cur.nxt[i] = new TrieNode();
      cur = cur.nxt[i]!;
      cur.cnt++;
    }
    cur.end++;
  }
  countWordsEqualTo(w: string): number {
    let cur = this.root;
    for (const c of w) {
      const i = c.charCodeAt(0) - 97;
      if (!cur.nxt[i]) return 0;
      cur = cur.nxt[i]!;
    }
    return cur.end;
  }
  countWordsStartingWith(p: string): number {
    let cur = this.root;
    for (const c of p) {
      const i = c.charCodeAt(0) - 97;
      if (!cur.nxt[i]) return 0;
      cur = cur.nxt[i]!;
    }
    return cur.cnt;
  }
  erase(w: string): void {
    let cur = this.root;
    for (const c of w) {
      const i = c.charCodeAt(0) - 97;
      cur = cur.nxt[i]!;
      cur.cnt--;
    }
    cur.end--;
  }
}

Complexity

• ⏰ Time complexity: O(l) — Each operation traverses the word/prefix of length l.
• 🧺 Space complexity: O(n * l) — n is the number of words, l is the average word length (for Trie storage).