Untitled

class Search:
    def __init__(self, word_counts):
        # init trie
        self.trie = {}
        
        # ingest every  pair
        for word in word_counts:
            self.insert(word, word_counts[word])
    
    def insert(self, word, count):
        # start at root node
        node = self.trie
        
        # iterate through chars in word
        for char in word:
            if char not in node:
                node[char] = {}
            node = node[char]
    
        # at the end of the word, point to payload
        node["word"] = word
        node["count"] = count
        
    def walk_to_node(self, prefix):
        # start at root node
        node = self.trie
        
        # walk to node at end of prefix
        for char in prefix:
            if char not in node:
                raise Exception
            else:
                node = node[char]
                
        # Return trie node
        return node
        
    def search(self, node):
        # stack
        s = []
        
        # seed with node
        s.append(node)
        
        # container for output
        count_pairs = []
        
        # run search
        while s:
            # get cur
            cur = s.pop()
            
            # check if cur is a leaf, otherwise, continue
            # down tree of children
            if "word" in cur:
                count_pairs.append((cur['word'], cur['count']))
            for child in cur:
                if child not in ('word', 'count'):
                    s.append(cur[child])
                    
        # return count pairs (truncated by search results)
        return count_pairs
    
    def suggest(self, prefix, num_results):
        # find the trie node to initiate search
        trie_node =self.walk_to_node(prefix)
        
        # run search from this node and return
        search_results = self.search(trie_node)
        
        # sort search results by the count
        search_results.sort(key = lambda pair: -1 * pair[1])
        
        # return, truncating by number of results desired
        return search_results[:num_results]
    
word_counts = {
    "apple": 150,
    "application": 100,
    "apricot": 75,
    "banana": 200,
    "bandana": 50,
    "baddadan": 100,
    "cantaloupe": 60,
    "cat": 120,
    "caterpillar": 80,
    "pea": 20,
    "peach": 135,
    "pear": 145,
    "pineapple": 160,
    "plum": 155
}

search = Search(word_counts)
print(search.suggest('', num_results = 10))
print(search.suggest('app', num_results = 10))
print(search.suggest('pea', num_results = 10))
print(search.suggest('b', num_results = 2))
print(search.suggest('pea', num_results = 2))
search.insert('zoolander', 10)
print(search.suggest('zoo', num_results = 2))