Introduction
A trie is just a tree representing words. Each node represents a character. In its most fundamental form, it supports the following functions: insert(word), search(word).
A lot of difficult problems can be solved/optimized by:
1. Inserting words into a Trie
2. Traversing the trie to find all words that satisfy some property
Trie's are useful because they are time/space efficient. Words that share a prefix are not duplicated in the trie - they use the same prefix path.
This is what the trie looks like after inserting the 3 words: "top", "bear", "beer". (order doesn't matter)
root
/ \
t b
| |
o e
| |
p / \
a e
| |
r r
(bear) (beer)We see that since 'bear' and 'beer' share the same prefix 'be', we can represent both words in 6 nodes instead of 8.
O(len(word)) for insert and searchO(num_words * avg_word_length) in worst case of no overlapping prefixesTrie Template
The Trie() class represents a node in the trie. Each node contains a map to its children nodes of type: {char -> Trie}. It also contains a boolean, indicating whether this node is the last character of a word or not.
This is set during insertion, and necessary when searching the trie to tell if this node is a word or not. In the example trie above, the 2 'r' nodes and 'p' nodes have this boolean true, rest nodes are false.
You may be wondering, aren't leaf nodes guaranteed to be words, why do we need to store this boolean? It is true that leaf nodes are words, but interior nodes can be words as well. Imagine we insert 'bears' into the above trie. 'Bear' is a word, but the last 'r' in 'bear' is an interior node of the Trie.
Basic Template
Insert: We initialize the current node as the current Trie. We iterate over the word. If the character is not in the trie (children of the root node) we create it in the map. We traverse to that node, and repeat. Finally, we set the last node in the word to True.
Search: Same traversal, except if the character is not in the trie, we return False. And after we are at the final node, we need to return if the node is at the end of a word or not.
For example: search('beer') = True, search('bee') = False in the above Trie.
It's important to note that this is just a template, and you should be open to modifying the logic as needed to solve the problem.
class Trie:
def __init__(self):
self.is_word = False
self.children = {}
def insert(self, word: str) -> None:
node = self
for c in word:
if c not in node.children: node.children[c] = Trie()
node = node.children[c]
node.is_word = True
def search(self, word: str) -> bool:
node = self
for c in word:
if c not in node.children: return False
node = node.children[c]
return node.is_wordBasic Trie Problems
Classic problems implementing and extending the basic trie structure.
01Implement Trie (Prefix Tree)Medium
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.
def startsWith(self, prefix: str) -> bool:
node = self
for c in prefix:
if c not in node.children: return False
node = node.children[c]
return TrueO(len(word))O(total characters)It's the same template above, just with a startsWith function. Its the same traversal logic as search, except at the end we just return True instead of node.is_word, since we don't care if its a word or not. We just care that this prefix exists in the trie, meaning some word that was inserted has this prefix.
02Implement Trie II (Prefix Tree)Medium
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.
class Trie:
def __init__(self):
self.children = {}
self.count = 0 # how many words go through this node
self.word_count = 0 # how many words end at this node
def insert(self, word: str) -> None:
node = self
for c in word:
if c not in node.children:
node.children[c] = Trie()
node = node.children[c]
node.count += 1
node.word_count += 1
def countWordsEqualTo(self, word: str) -> int:
node = self
for c in word:
if c not in node.children:
return 0
node = node.children[c]
return node.word_count
def countWordsStartingWith(self, prefix: str) -> int:
node = self
for c in prefix:
if c not in node.children:
return 0
node = node.children[c]
return node.count
def erase(self, word: str) -> None:
node = self
for c in word:
if c not in node.children:
return
node = node.children[c]
node.count -= 1
node.word_count -= 1O(len(word))O(total characters)This is the same idea. The only difference is we need to store a count variable representing how many words go through this node (how many words contain this node as a prefix) and word_count (how many words end at this node) at each node instead of just is_word.
It's interesting to note that for erase, ideally we can remove chars from the children map when count == 0 for proper memory management, but not strictly necessary for correctness. It would also be more robust, in the case we call erase(word) when word wasn't inserted. (in the constraints of this problem, we are guaranteed this won't happen, but useful to think about as a followup)
Search Variations
Problems that modify the search function to handle wildcards or edits.
01Design Add and Search Words Data StructureMedium
Design a data structure that supports adding new words and finding if a string matches any previously added string.
Implement the WordDictionary class:
- WordDictionary() Initializes the object.
- void addWord(word) Adds word to the data structure, it can be matched later.
- bool search(word) Returns true if there is any string in the data structure that matches word or false otherwise. word may contain dots '.' where dots can be matched with any letter.
def search(self, word: str) -> bool:
node = self
for i in range(len(word)):
c = word[i]
if c == '.':
return any(child.search(word[i+1:]) for child in node.children.values())
if c not in node.children: return False
node = node.children[c]
return node.is_wordO(26^m) worst case where m is number of dotsO(1)Standard template, but the only tricky part is handling periods. We need to recursively call search for word[i+1:] when we see a period, on all the children of the current node, as we can match onto any of them.
02Words Within Two Edits of DictionaryMedium
You are given two string arrays, queries and dictionary. All words in each array comprise of lowercase English letters and have the same length.
In one edit you can take a word from queries, and change any letter in it to any other letter. Find all words from queries that, after a maximum of two edits, equal some word from dictionary.
Return a list of all words from queries, that match with some word from dictionary after a maximum of two edits. Return the words in the same order they appear in queries.
Brute Force Solution
def twoEditWords(self, queries: List[str], dictionary: List[str]) -> List[str]:
res = []
n = len(dictionary[0])
for query in queries:
for word in dictionary:
if sum(query[i] != word[i] for i in range(n)) <= 2:
res.append(query)
break # NEED THE BREAK
return res
# 1 liner
def twoEditWords(self, queries: List[str], dictionary: List[str]) -> List[str]:
return [query for query in queries if any(sum(c1 != c2 for c1, c2 in zip(query, word)) <= 2 for word in dictionary)]Trie Solution
class Trie:
def __init__(self):
self.is_word = False
self.children = defaultdict(Trie)
def insert(self, word: str) -> None:
node = self
for c in word:
node = node.children[c] # only 1 line needed!
node.is_word = True
def search(self, word, i, allowable_mismatch_rem) -> bool:
if allowable_mismatch_rem < 0: return False
if i == len(word): return True
res = False
for c, trie in self.children.items():
res |= trie.search(word, i+1, allowable_mismatch_rem - (c != word[i]))
return res
class Solution:
def twoEditWords(self, queries: List[str], dictionary: List[str]) -> List[str]:
trie = Trie()
for word in dictionary:
trie.insert(word)
return [query for query in queries if trie.search(query, 0, 2)]O(n * m * 26^2)O(total characters in dictionary)This is a simple trie application. We try all children in the trie, and when there is a mismatch while traversing the trie, we reduce the edit count by one. When the edit count is negative we return False. If we reach the end of the word with at most 2 edits, this word works. If any such children work, we return True.
XOR Trie
A specialized trie for solving XOR-related problems efficiently.
Basic XOR (^) properties:
- 1 XOR 0 = 0 XOR 1 = 1
- 1 XOR 1 = 0 XOR 0 = 0
- x ^ x = 0
- x ^ 0 = x
The inverse of XOR operator is the XOR operator! ie. x XOR y = z. We can XOR both sides by x: x XOR (x XOR y) = x XOR z → (x XOR x) XOR y = x XOR z → 0 XOR y = x XOR z → y = x XOR z. We can do this to isolate for y.
Consider finding the maximum XOR between any two numbers in an array. The brute force approach would compare every pair - O(n²). However, by viewing numbers in their binary representation and building a trie, we can find the maximum XOR partner for each number in O(32) = O(1) time, leading to an O(n) solution.
An XOR trie stores binary representations vertically, with each level representing one bit position. The XOR trie is always 32 levels deep.
XOR Trie Template
For each number, we want to find another number that differs in as many significant bit positions as possible, maximizing XOR. Process bits from most significant (31) to least (0).
class TrieNode:
def __init__(self):
self.children = {} # Maps 0/1 to child nodes
self.count = 0 # Number of numbers ending here (optional)
class XORTrie:
def __init__(self):
self.root = TrieNode()
def insert(self, num: int) -> None:
node = self.root
# Process each bit from most significant (31) to least (0)
for i in range(31, -1, -1):
bit = (num >> i) & 1
if bit not in node.children:
node.children[bit] = TrieNode()
node = node.children[bit]
node.count += 1
def findMaxXOR(self, num: int) -> int:
if not self.root.children:
return 0
node = self.root
result = 0
# Try to go opposite direction at each bit
for i in range(31, -1, -1):
bit = (num >> i) & 1
opposite = 1 - bit
# Take opposite path if exists (maximizes XOR)
if opposite in node.children:
result |= (1 << i) # Set this bit in result
node = node.children[opposite]
else:
node = node.children[bit]
return result01Maximum XOR of Two Numbers in an ArrayMedium
Given an integer array nums, return the maximum result of nums[i] XOR nums[j], where 0 ≤ i, j < n.
Example:
Input: nums = [3,10,5,25,2,8]
Output: 28
Explanation: The maximum XOR is achieved by 5 XOR 25 = 28
def findMaximumXOR(self, nums: List[int]) -> int:
trie = XORTrie()
for num in nums:
trie.insert(num)
max_xor = 0
for num in nums:
max_xor = max(max_xor, trie.findMaxXOR(num))
return max_xorO(n)O(n)For each number, we want to find another number that differs in as many significant bit positions as possible, maximizing XOR.
02Maximum XOR With an Element From ArrayHard
Given an array nums and queries of the form [xi, mi], find the maximum XOR value of xi with any number in nums that is less than or equal to mi. Return -1 if no number in nums is less than or equal to mi.
class Trie:
def __init__(self):
self.children = {}
def insert(self, num: int) -> None:
node = self
for i in range(31, -1, -1):
c = (num >> i) & 1
if c not in node.children: node.children[c] = Trie()
node = node.children[c]
def query(self, num: int) -> int:
if not self.children: return -1 # NEED THIS IN CASE 0 INSERT() were called... ie. trie is empty
node = self
res = 0
for i in range(31, -1, -1):
c = (num >> i) & 1
if 1 - c in node.children: # going here will max the XOR
node = node.children[1-c]
res |= (1 << i)
else: # if there is at least one num in the trie, we are guaranteed either 1-c or c is in the children
node = node.children[c]
return res
class Solution:
# sort both nums,queries. then two ptrs.
# O(nlogn + mlogm)
def maximizeXor(self, nums: List[int], queries: List[List[int]]) -> List[int]:
xor_trie = Trie()
nums.sort()
j = 0
n = len(nums)
queries = sorted(enumerate(queries), key=lambda x: x[1][1])
res = [-1] * len(queries)
for i, (x, m) in queries:
while j < n and nums[j] <= m:
xor_trie.insert(nums[j])
j += 1
if j > 0: # need some value in the trie... if all nums[j] > m, answer is -1.
res[i] = xor_trie.query(x)
return resO(n log n + m log m)O(n)This requires a trick where BECAUSE we have the queries all at once (ie, not ONLINE) we can sort them and answer them in an order most advantageous for us. We can sort the queries based on the limit value, and only add numbers in the array up to that limit to satisfy the requirement in the XOR_TRIE. Now, we just simply use the xor_trie, and re-wire the indices of the query into its original place.
Summary
Key Trie Concepts:
1. Each node represents a character
2. Paths from root represent prefixes/words
3. is_word boolean marks word endings
4. Children map: {char -> TrieNode}
Standard Trie Applications:
- Autocomplete
- Spell checking
- Prefix matching
- Word search with wildcards
XOR Tries excel at:
1. Finding maximum/minimum XOR pairs
2. Counting XOR pairs within ranges
3. Dynamic XOR queries on changing sets
4. Subarray XOR optimizations
Remember:
- Always process bits from most to least significant in XOR tries
- Consider using prefix XORs for subarray problems
- Augment trie nodes with additional info (min values, counts) as needed
- Use DFS for tree-based XOR problems