Implement a Prefix Tree (Trie)

Apr 4, 2019 by Kevin Nasto

Description

Implement a trie (prefix tree) with insert, search, and startsWith methods. Tries are used in many applications such as spell checking, autocompletion, and ip routing.

Part 2 of this problem (not shown on LeetCode) is to implement an autoComplete method.

Solution 1: Using Node Class and Storing Entire Word at Leaf

This solution uses a Node class which makes things easier to keep track of. Storing the entire word in the end of word seperator makes the implementation less tricky. However this unfortunetly causes the space complexity to be O(n) where n is the number of words.

class Node:

    def __init__(self, letter):
        self.letter = letter
        self.children = {}
        self.end = False

class Trie:

    def __init__(self):
        self.root = Node('root')

    def insert(self, word):
        node = self.root
        for letter in word:
            if letter not in node.children:
                node.children[letter] = Node(letter)
            node = node.children[letter]
        node.end = word

    def search(self, word):
        node = self.root
        for letter in word:
            if letter in node.children:
                node = node.children[letter]
            else:
                return False
        if node.end:
            return True
        else:
            return False

    def startsWith(self, prefix):
        node = self.root
        for letter in prefix:
            if letter in node.children:
                node = node.children[letter]
            else:
                return False
        return True

    def autocomplete(self, word):
        node = self.root
        for letter in word:
            if letter in node.children:
                node = node.children[letter]
            else:
                return []
        words = []
        autocomplete_helper(node, words)
        return words


def autocomplete_helper(node, words):
    if node.end:
        words.append(node.end)
        return
    for child in node.children:
        autocomplete_helper(node, words)

Solution 2: Using Just Dictionary and End of Word Marker $

Since this only stores the letters, and not the entire word, the space complexity is much less.

class Trie:

    def __init__(self):
       self.root = {}

    def insert(self, word):
        node = self.root
        for letter in word:
            if letter not in node:
                node[letter] = {}
            node = node[letter]
        node['$'] = {} # End of word

    def search(self, word):
        node = self.root
        for letter in word:
            if letter in node:
                node = node[letter]
            else:
                return False
        if '$' in node:
            return True
        else:
            return False

    def startsWith(self, prefix):
        node = self.root
        for letter in prefix:
            if letter in node:
                node = node[letter]
            else:
                return False
        return True

    def autoComplete(self, prefix):
        node = self.root
        for letter in prefix:
            if letter in node:
                node = node[letter]
            else:
                return []
        words = []
        autocomplete_helper(prefix, words, node)
        return words


def autocomplete_helper(parent, words, node):
    if '$' in node:
        words.append(parent)
    for letter in node:
        print_tree(parent + letter, words, node[letter])

Test Code and Benchmarking

trie = Trie()
trie.insert('cat')
trie.insert('can')
trie.insert('and')
trie.insert('ant')
trie.insert('antenna')

print(trie.autocomplete('an'))

Using a list of 8000 most frequently searched words, here are some benchmarks.

77848 Bytes using python list
1184  Bytes using above trie

Input: p
Trie:  4.71 ms
List:  4.28 ms

Input: cat
Trie: 0.25 ms
List: 5.19 ms

Input: python
Trie:  0.05 ms
List:  7.81 ms