class TrieNode:
def __init__(self):
self.children = {}
self.topSearch = []
class Trie:
def __init__(self):
self.root = TrieNode()
self.hm = defaultdict(int)
def insert(self, query, times):
self.hm[query] += times
cur = self.root
for c in query:
if c not in cur.children:
cur.children[c] = TrieNode()
cur = cur.children[c]
# Add the query in topSearch if it matches the given condition
if query not in cur.topSearch: cur.topSearch.append(query)
cur.topSearch.sort(key = lambda x: (-self.hm[x], x))
if len(cur.topSearch) > 3:
del cur.topSearch[-1]
class AutocompleteSystem:
def __init__(self, sentences: List[str], times: List[int]):
self.trie = Trie()
self.query = '' # to keep a record of all seen characters in this query
self.cur = self.trie.root # to directly go to the children of current TrieNode
self.lastQueryResult = True # to track if we have seen all the characters from the query in our pre-stored trie
for i in range(len(sentences)):
self.trie.insert(sentences[i], times[i])
def input(self, c: str) -> List[str]:
if c == '#':
self.trie.insert(self.query, 1)
self.query = ''
self.cur = self.trie.root
self.lastQueryResult = True
return []
self.query += c
if not self.lastQueryResult:
return []
if c not in self.cur.children:
self.lastQueryResult = False
return []
# move to the next child and return topSearch
self.cur = self.cur.children[c]
return self.cur.topSearch