Trie (Prefix Tree) Implementation in Python
A Trie, also known as a prefix tree, is a tree-like data structure that is used to store a dynamic set of strings where the keys are usually strings. It is commonly used for tasks like autocomplete, spell checking, and IP routing due to its efficient retrieval of data.
Program Structure
The Trie implementation in Python consists of the following components:
1. TrieNode Class
This class represents a single node in the Trie. It includes the following attributes:
children
: A dictionary mapping each character to its corresponding child TrieNode.is_end_of_word
: A boolean indicating whether the node marks the end of a valid word.
2. Trie Class
This class encapsulates the functionality of the Trie. It includes the following methods:
__init__(self)
: Initializes the Trie with an empty root node.insert(self, word)
: Inserts a word into the Trie.search(self, word)
: Searches for a word in the Trie.starts_with(self, prefix)
: Checks if there is any word in the Trie that starts with a given prefix.
Python Code Implementation
class TrieNode:
def __init__(self):
"""
Initialize a TrieNode.
"""
self.children = {} # Maps each character to the corresponding TrieNode
self.is_end_of_word = False # True if the node represents the end of a word
class Trie:
def __init__(self):
"""
Initialize the Trie.
"""
self.root = TrieNode()
def insert(self, word):
"""
Insert a word into the Trie.
:param word: The word to insert.
"""
node = self.root
for char in word:
if char not in node.children:
node.children[char] = TrieNode()
node = node.children[char]
node.is_end_of_word = True
def search(self, word):
"""
Search for a word in the Trie.
:param word: The word to search for.
:return: True if the word is found, False otherwise.
"""
node = self.root
for char in word:
if char not in node.children:
return False
node = node.children[char]
return node.is_end_of_word
def starts_with(self, prefix):
"""
Check if there is any word in the Trie that starts with the given prefix.
:param prefix: The prefix to check.
:return: True if there is any word with the given prefix, False otherwise.
"""
node = self.root
for char in prefix:
if char not in node.children:
return False
node = node.children[char]
return True
Explanation
The Trie implementation is designed to efficiently store and retrieve strings, particularly for operations that involve searching for words or prefixes. Here’s how the implementation works:
1. TrieNode Class
The TrieNode
class represents a single node in the Trie. Each node contains:
children
: A dictionary that maps characters to their corresponding child TrieNodes.is_end_of_word
: A boolean flag that indicates whether the node represents the end of a valid word in the Trie.
2. Trie Class
The Trie
class manages the operations of the Trie.
Insertion
The insert
method adds a word to the Trie. It starts at the root node and traverses through the Trie by following the characters of the word. If a character does not have a corresponding child node, a new TrieNode
is created. Once the entire word is inserted, the last node is marked as the end of the word.
Search
The search
method checks whether a word exists in the Trie. It traverses the Trie using the characters of the word. If it reaches a node where a character is not found, it returns False
. If it successfully traverses all characters and the last node is marked as the end of a word, it returns True
.
Prefix Check
The starts_with
method checks if there is any word in the Trie that starts with a given prefix. It follows the prefix characters in the Trie, and if it successfully finds all the characters, it returns True
.
Usage Example
# Example usage of the Trie
# Initialize the Trie and insert words
trie = Trie()
trie.insert("apple")
trie.insert("app")
trie.insert("bat")
trie.insert("ball")
# Search for words
print("Search 'apple':", trie.search("apple")) # Output: True
print("Search 'app':", trie.search("app")) # Output: True
print("Search 'batman':", trie.search("batman")) # Output: False
# Check for prefixes
print("Starts with 'ba':", trie.starts_with("ba")) # Output: True
print("Starts with 'bat':", trie.starts_with("bat")) # Output: True
print("Starts with 'cat':", trie.starts_with("cat")) # Output: False
This example demonstrates how to create a Trie, insert words into it, and perform both word searches and prefix checks.