Data Structures

Implement a Trie

Write a Class designed for storing and looking up strings. The class should have a method for adding a new string add, checking if a string was aready added has, and checking if a string starts with some prefix hasPrefix.

Example:

x = Strings()

x.add("pizza")
x.add("zap")
x.add("pie")

x.has("pizza") # True
x.has("pi")    # False

x.hasPrefix("pi")  # True
x.hasPrefix("pa")  # False

c.has('pie') # False
c.hasPrefix('pie') # False
c.hasPrefix('piz') # False

First Test Case:

c = Strings()
c.add('pizza')
c.add('zap')
c.add('pie')
c.has('pizza')
c.has('pi')
c.hasPrefix('pi')
c.hasPrefix('pa')
c.has('pie')
c.hasPrefix('pie')
c.hasPrefix('piz')
# We compare the result of every function call when judging.

We're going to be implementing a Trie to solve this problem, since Tries give us fast prefix lookups. See the previous lesson for details, but the main idea is to turn the strings we've seen into a Trie like this:

We saw that one easy way to create a Trie is to nest many nested HashMaps. This is how we'll implement our Trie:

trie = {
    "p": {
        "i": {
            "z": {
                "z": {
                    "a": {}
                }
            },
            "e": {}
        }
    },
    "z": {
        "a": {
            "p": {}
        }
    }
}

# traverse the Trie like this:
trie["p"]["i"]["e"]

To add a string to the Trie, we can walk letter by letter in the string and making sure there's a node for each letter. If there's not a node, we'll add one in.

def add(self, string):
    node = self.root  # the node we're standing at
    
    # walk letter by letter, adding nodes if they don't show up
    for l in string: 
        if l not in node:
            node[l] = {}
        node = node[l]
    node['isEndOfWord'] = True # store whether this is an ending letter of a word

To find if a prefix is in the Trie, we can walk letter-by-letter, making sure there's a node for each letter.

def hasPrefix(self, prefix):
    node = self.root  # the node we're standing at

    # walk letter by letter. Return False if there's no node to walk to.
    for l in prefix:
        if l not in node:
            return False
        node = node[l]
    return True

To find if an entire string is in the Trie, we can use the same exact code as above. The only difference is we have to check whether the node we end up on is the end to one of the words we've seen. Here's the code for this:

def has(self, string):
    node = self.root  # the node we're standing at

    # walk letter by letter
    for l in string:
        if l not in node:
            return False
        node = node[l]
    return 'isEndOfWord' in node # return whether we get to the end of a word

To code this up, we can combine all of these methods. The initialization function needs to set up the Trie right by storing the root node of the tree root (an empty HashMap at the start). Here's the full code:

Open SpoilerClose Spoiler

class Strings:

    def __init__(self):
        self.root = {} # the root node of the Trie


    def add(self, string):
        node = self.root # the node we're standing at
        for l in string:
            if l not in node:
                node[l] = {}
            node = node[l]
        node['isEndOfWord'] = True


    def hasPrefix(self, prefix):
        node = self.root # the node we're standing at
        for l in prefix:
            if l not in node:
                return False
            node = node[l]
        return True


    def has(self, string):
        node = self.root # the node we're standing at
        for l in string:
            if l not in node:
                return False
            node = node[l]
        return 'isEndOfWord' in node

add Complexity: $O(m)$ , where $m$ is the length of the inserted string.

hasPrefix Complexity: $O(m)$

has Complexity: $O(1)$

init Complexity: $O(1)$

Mark as Completed:

Submits:

Strings

Test your code to get an output here!