Skip to content

Trie Implementation

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end_of_word = False


class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        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):
        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):
        node = self.root
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        return True


trie = Trie()
words = ["apple", "banana", "orange", "pear"]
for word in words:
    trie.insert(word)

print(trie.search("banana"))  # True
print(trie.search("grape"))  # False
print(trie.starts_with("or"))  # True
print(trie.starts_with("an"))  # False