Search Suggestions System

search_suggest_system.py

"""

- given: products, searchword
- suggests at most 3 (common prefix)
    - after each character is typed

BF: (P.W + P log P) * S
- search arr, get all words with prefix
- sort by lexographic order, return top 3

SOL 1:
- we can store words in trie

P*S
get_words(pref, trie): 
    - heap
    - for each word, add it to a heap, of max size 3
    - return heap

Input
["havana"]
"tatiana"
Output:
[]
Expected
[[],[],[],[],[],[],[]]


"""
import heapq

class Word:
    def __init__(self, word):
        self.word = word

    def __gt__(self, other):
        return other.word > self.word

class Trie:
    def __init__(self):
        self.store = {}
        self.end_char = "*"

    def add(self, word):
        curr_store = self.store

        for char in word:
            if char not in curr_store:
                curr_store[char] = {}
            curr_store = curr_store[char]

        curr_store[self.end_char] = word

    def search_prefix(self, prefix):
        curr_store = self.store

        for char in prefix:
            if char not in curr_store:
                return []
            curr_store = curr_store[char]

        res = []
        self._get_all_words(curr_store, res)
        return res
        
    def _get_all_words(self, curr_store, result):

        for key, value in curr_store.items():
            if key == self.end_char:
                heapq.heappush(result, Word(value))
                if len(result) > 3:
                    heapq.heappop(result)
            else:
                self._get_all_words(value, result)


class Solution:
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        
        result = []

        trie = Trie()
        for word in products:
            trie.add(word)

        for i in range(len(searchWord)):
            search_results = trie.search_prefix(searchWord[:i+1])
            # print(i, searchWord[:i+1], search_results)
            if search_results:
                res = [item.word for item in search_results]
                res.sort()
                result.append(res)
            else:
                result.append([])

            # print(result)

        return result