Implement Trie (Prefix Tree)

MediumTrieDesignString
Category: Tries
Companies that ask this question:
AmazonGoogleFacebookMicrosoftUber

Approach

Implement Trie (Prefix Tree)

Problem Statement

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.
  • boolean search(String word) Returns true if the string word is in the trie (i.e., was inserted before), and false otherwise.
  • boolean startsWith(String prefix) Returns true if there is a previously inserted string word that has the prefix prefix, and false otherwise.

Examples

Example 1:

Input:
["Trie", "insert", "search", "search", "startsWith", "insert", "search"]
[[], ["apple"], ["apple"], ["app"], ["app"], ["app"], ["app"]]

Output:
[null, null, true, false, true, null, true]

Explanation:
Trie trie = new Trie();
trie.insert("apple");
trie.search("apple");   // return True
trie.search("app");     // return False
trie.startsWith("app"); // return True
trie.insert("app");
trie.search("app");     // return True

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, search, and startsWith.

Approach

Key Insight

A Trie is a tree where:

  • Each node represents a character
  • Paths from root represent words/prefixes
  • Each node has up to 26 children (a-z)
  • Mark end of word with a flag

Trie Node Structure

class TrieNode:
    def __init__(self):
        self.children = {}  # or [None] * 26
        self.is_end_of_word = False

Operations

Insert (word)

  1. Start at root
  2. For each character in word:
    • If child doesn't exist, create it
    • Move to child node
  3. Mark last node as end of word

Search (word)

  1. Start at root
  2. For each character:
    • If child doesn't exist, return False
    • Move to child
  3. Return True only if last node is marked as end of word

StartsWith (prefix)

  1. Start at root
  2. For each character:
    • If child doesn't exist, return False
    • Move to child
  3. Return True (don't need end of word check)

Complexity Analysis

Time Complexity:

  • Insert: O(m) where m = word length
  • Search: O(m) where m = word length
  • StartsWith: O(m) where m = prefix length

Space Complexity:

  • Overall: O(ALPHABET_SIZE × N × M)
    • N = number of words
    • M = average word length
    • ALPHABET_SIZE = 26

Pattern Recognition

This problem demonstrates:

  • Trie data structure fundamentals
  • Prefix matching efficiently
  • Design pattern questions
  • Foundation for autocomplete, spell checker
  • Used in word search problems

Solution

java
1import java.util.*;
2
3class Trie {
4    class TrieNode {
5        Map<Character, TrieNode> children;
6        boolean isEndOfWord;
7
8        TrieNode() {
9            children = new HashMap<>();
10            isEndOfWord = false;
11        }
12    }
13
14    private TrieNode root;
15
16    public Trie() {
17        root = new TrieNode();
18    }
19
20    public void insert(String word) {
21        TrieNode node = root;
22
23        for (char ch : word.toCharArray()) {
24            node.children.putIfAbsent(ch, new TrieNode());
25            node = node.children.get(ch);
26        }
27
28        node.isEndOfWord = true;
29    }
30
31    public boolean search(String word) {
32        TrieNode node = root;
33
34        for (char ch : word.toCharArray()) {
35            if (!node.children.containsKey(ch)) {
36                return false;
37            }
38            node = node.children.get(ch);
39        }
40
41        return node.isEndOfWord;
42    }
43
44    public boolean startsWith(String prefix) {
45        TrieNode node = root;
46
47        for (char ch : prefix.toCharArray()) {
48            if (!node.children.containsKey(ch)) {
49                return false;
50            }
51            node = node.children.get(ch);
52        }
53
54        return true;
55    }
56}
57
58// Alternative implementation using array
59class TrieArray {
60    class TrieNode {
61        TrieNode[] children;
62        boolean isEndOfWord;
63
64        TrieNode() {
65            children = new TrieNode[26];  // a-z
66            isEndOfWord = false;
67        }
68    }
69
70    private TrieNode root;
71
72    public TrieArray() {
73        root = new TrieNode();
74    }
75
76    private int charToIndex(char ch) {
77        return ch - 'a';
78    }
79
80    public void insert(String word) {
81        TrieNode node = root;
82
83        for (char ch : word.toCharArray()) {
84            int index = charToIndex(ch);
85            if (node.children[index] == null) {
86                node.children[index] = new TrieNode();
87            }
88            node = node.children[index];
89        }
90
91        node.isEndOfWord = true;
92    }
93
94    public boolean search(String word) {
95        TrieNode node = root;
96
97        for (char ch : word.toCharArray()) {
98            int index = charToIndex(ch);
99            if (node.children[index] == null) {
100                return false;
101            }
102            node = node.children[index];
103        }
104
105        return node.isEndOfWord;
106    }
107
108    public boolean startsWith(String prefix) {
109        TrieNode node = root;
110
111        for (char ch : prefix.toCharArray()) {
112            int index = charToIndex(ch);
113            if (node.children[index] == null) {
114                return false;
115            }
116            node = node.children[index];
117        }
118
119        return true;
120    }
121}
122
123/**
124 * Your Trie object will be instantiated and called as such:
125 * Trie obj = new Trie();
126 * obj.insert(word);
127 * boolean param_2 = obj.search(word);
128 * boolean param_3 = obj.startsWith(prefix);
129 */
130
Loading visualizer...