How to return an int within a trie in python

Viewed 57

I'm trying to search for a word within my trie data structure and if it exists in the trie, return the word and its frequency.

This is what the trie looks like:

{'c': {'u': {'t': {('e', 10): True, ('s', 50): True},
   ('t', 30): True,
   ('b', 15): True},
  'a': {'l': {('m', 1000): True}}},
 'a': {'n': {('t', 20): True,
   'n': {'o': {'t': {'a': {'t': {'i': {'o': {('n', 5): True}}}}}}}},
  'p': {'p': {'l': {('e', 300): True},
    'e': {'n': {'d': {'i': {('x', 10): True}}}},
    ('s', 60): True},
   'o': {'l': {'o': {'g': {('y', 600): True,
       'e': {'t': {'i': {('c', 1000): True}}}}}}}}},
 'f': {'u': {'r': {'t': {'h': {'e': {('r', 40): True}}},
    'n': {'i': {'t': {'u': {'r': {('e', 500): True}}}}}},
   ('r', 10): True},
  'i': {'n': {('d', 400): True}},
  'a': {'r': {('m', 5000): True,
    'm': {'i': {'n': {('g', 1000): True}}, 'e': {('r', 300): True}}},
   't': {'h': {'o': {('m', 40): True}}}}}}

Say I'm searching for the word cuts, how would I go about coding that so that I can return its integer pair which in this case is 50?

Btw the True part marks the end of the word.

This is my search function in psedocode format:

def search(self, word: str) -> int:

    current = self.root

    if word in current.children:
        return its integer pair if its the end of a word and matches
            
    return 0

Here's the Trie class:

# Class representing a node in the Trie

class TrieNode:

def __init__(self, letter=None, frequency=None, is_last=False):
    self.letter = letter            # letter stored at this node
    self.frequency = frequency      # frequency of the word if this letter is the end of a word
    self.is_last = is_last          # True if this letter is the end of a word
    self.children: dict[str, TrieNode] = {}     # a hashtable containing children nodes, key = letter, value = child node

class TrieDictionary(BaseDictionary):

def __init__(self):
    # TO BE IMPLEMENTED
    #setting root to object TrieNode()
    self.root = TrieNode()

Basically for the output of my trie I printed out self.root.children

2 Answers

You can do this recursively:

def recursive_tree_search(tree, word):
    if len(word) == 1:
        return {k: v for k, v in tree.keys()}[word]
    else:
        subtree = tree[word[0]]
        subword = word[1:]
        return recursive_tree_search(subtree, subword)


example_tree = {
    'c': {
        'u': {
            't': {
                ('e', 10): True, ('s', 50): True}, ('t', 30): True, ('b', 15): True},
        'a': {
            'l': {
                ('m', 1000): True}}},
    'a': {
        'n': {
            ('t', 20): True,
            'n': {
                'o': {
                    't': {
                        'a': {
                            't': {
                                'i': {
                                    'o': {
                                        ('n', 5): True}}}}}}}},
        'p': {
            'p': {
                'l': {
                    ('e', 300): True},
                'e': {
                    'n': {
                        'd': {
                            'i': {
                                ('x', 10): True}}}},
                ('s', 60): True},
            'o': {
                'l': {
                    'o': {
                        'g': {
                            ('y', 600): True,
                            'e': {
                                't': {
                                    'i': {
                                        ('c', 1000): True}}}}}}}}},
    'f': {
        'u': {
            'r': {
                't': {
                    'h': {
                        'e': {
                            ('r', 40): True}}},
                'n': {
                    'i': {
                        't': {
                            'u': {
                                'r': {
                                    ('e', 500): True}}}}}},
            ('r', 10): True},
        'i': {
            'n': {
                ('d', 400): True}},
        'a': {
            'r': {
                ('m', 5000): True,
                'm': {
                    'i': {
                        'n': {
                            ('g', 1000): True}},
                    'e': {
                        ('r', 300): True}}},
            't': {
                'h': {
                    'o': {
                        ('m', 40): True}}}}}}


print(recursive_tree_search(example_tree, "cuts"))

Result: 50

Note: This doesn't have any Error-handling if a key isn't found and things like that. But that should be easy to add.

Trie-Node-class-based version:

def search(self, word: str) -> int:
    if len(word) == 1:
        return self.frequency
    else:
        return self.children[word[0]].search(word[1:])

No guarantees here, as I can't test this. But something like that should work.

The dictionary you provide at the start of your question seems to be some output generated from the TrieDictionary class, as it doesn't correspond well to that class definition itself: it has no tuples as dictionary keys.

Here is how the search method could look:

def search(self, word: str) -> int:
    current = self.root
    for ch in word:
        if ch not in current.children:
            return 0  # not found
        current = current.children[ch]
    if current.is_last:
        return current.frequency
    return 0  # it's a proper prefix of a trie-word, but not matching one 

For completeness sake, I provide here the code that loads the data in the trie using an insert method, and then calls the search method to retrieve the frequency of the word "cuts":

class TrieNode:
    def __init__(self, letter=None, frequency=None, is_last=False):
        self.letter = letter
        self.frequency = frequency
        self.is_last = is_last
        self.children: dict[str, TrieNode] = {}

class TrieDictionary:
    def __init__(self):
        self.root = TrieNode()
    
    def insert(self, word, freq=1):
        node = self.root
        for ch in word:
            if not ch in node.children:
                node.children[ch] = TrieNode(ch)
            node = node.children[ch]
        if node.is_last:
            node.frequency += freq
        else:
            node.is_last = True
            node.frequency = freq

    def search(self, word: str) -> int:
        current = self.root
        for ch in word:
            if ch not in current.children:
                return 0  # not found
            current = current.children[ch]
        if current.is_last:
            return current.frequency
        return 0  # it's a proper prefix of a trie-word, but not matching one 

data = [
    ('cute', 10),
    ('cuts', 50),
    ('cut', 30),
    ('cub', 15),
    ('calm', 1000),
    ('ant', 20),
    ('annotation', 5),
    ('apple', 300),
    ('appendix', 10),
    ('apps', 60),
    ('apology', 600),
    ('apologetic', 1000),
    ('further', 40),
    ('furniture', 500),
    ('fur', 10),
    ('find', 400),
    ('farm', 5000),
    ('farming', 1000),
    ('farmer', 300),
    ('fathom', 40)
]

trie = TrieDictionary()
for word, freq in data:
    trie.insert(word, freq)

print(trie.search("cuts"))  # 50

Other remarks

Not your question, but the TrieNode class is bloated: it should not need to have a letter attribute, as that follows from the character that was selected in the trie to get to that node. It is redundant information. You can see in the proposed search implementation that this attribute is never used.

Also, TrieNode should not really need the is_last attribute: the fact that it is a terminating node is implied by a non-zero frequency.

That frequency attribute should always be an int (not None which currently is its default value). So either frequency has a value 0, and then the node is not a terminating node, or it is greater than 0, and then it is a terminating node.

Related