minte9

LearnRemember / ALGORITHMS

Insert

The trie is a kind of tree ideal for implementing autocomplete. This is how we can populate our trie with data.

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

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

    def insert(self, word):
        currentNode = self.root

        for c in word:

            # If the current node has child key with current character
            if currentNode.children.get(c):

                # Follow the child node:
                currentNode = currentNode.children[c]

            else:

                # Add the character as a new child node
                newNode = TrieNode()
                currentNode.children[c] = newNode

                # Follow this new node
                currentNode = newNode

T = Trie()
T.insert('ace')
T.insert('bad')
T.insert('act')

print(T.root.children)
print("a", T.root.children['a'].children)
print("a:c", T.root.children['a'].children['c'].children)

"""
    {'a', 'b': <__main__.TrieNode object at 0x7f80ec1ace80>}
    a {'c': <__main__.TrieNode object at 0x7f80ec1acee0>}
    a:c {'e': <__main__.TrieNode object at 0x7f80ec1f5e50>, 't': <__main__...}
"""

Search

We check if the current node has any children with the current character as the key. If there is such a child, we update the current node to be the child node.

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

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

    def insert(self, word):
        currentNode = self.root

        for c in word:
            if currentNode.children.get(c):
                currentNode = currentNode.children[c]
            else:
                newNode = TrieNode()
                currentNode.children[c] = newNode
                currentNode = newNode

        currentNode.children["*"] = None

    def search(self, word):
        currentNode = self.root

        for c in word:
            if currentNode.children.get(c):
                currentNode = currentNode.children.get(c)
            else:
                return None

        return currentNode

T = Trie()
T.insert('ace')
T.insert('bad')
T.insert('act')
T.insert('cat')
T.insert('bat')
T.insert('batter')

print(T.search('cat').children.items())
print(T.search('bat').children.items())

"""
    dict_items([('*', None)])
    dict_items([('*', None), ('t', <__main__.TrieNode...)])
"""

Autocomplete

We can implement our autocomplete feature.

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

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

    def insert(self, word):
        currentNode = self.root

        for c in word:
            if currentNode.children.get(c):
                currentNode = currentNode.children[c]
            else:
                newNode = TrieNode()
                currentNode.children[c] = newNode
                currentNode = newNode

        currentNode.children["*"] = None

    def search(self, word):
        currentNode = self.root

        for char in word:
            if currentNode.children.get(char):
                currentNode = currentNode.children.get(char)
            else:
                return None

        return currentNode

    def collectWords(self, node=None, prefix="", words=[]):
        currentNode = node or self.root
        
        for key, node in currentNode.children.items():
            if key == '*':
                words.append(prefix)
            else: 
                self.collectWords(node, prefix + key, words)

        return words

    def autocomplete(self, prefix):
        node = self.search(prefix)

        if not node: 
            return None

        return self.collectWords(node, prefix, words=[])

T = Trie()
for word in ['ace', 'bad', 'act', 'cat', 'bat', 'batter']:
    T.insert(word)

words = T.autocomplete('bat')
print("Autocomplete words for 'bat' = ", words)

words = T.autocomplete('ac')
print("Autocomplete words for 'ac' = ", words)

"""
    Autocomplete words for 'bat' =  ['bat', 'batter']
    Autocomplete words for 'ac' =  ['ace', 'act']
"""