class TrieNode:
def __init__(self):
self.children = {}
self.isWord = False
self.hot = 0
class Trie:
def __init__(self):
self.trie = TrieNode()
# tc O(M) M: length of sentence/word sc O(M)
def insert(self,word,freq):
d = self.trie
for c in word:
if c not in d.children:
d.children[c] = TrieNode()
d = d.children[c]
d.isWord = True
d.hot += freq
# tc O(M) sc O(number of nodes with prefix as node)
def search(self,target_str):
cur = self.trie
res = []
path = ""
for c in target_str:
if c not in cur.children:
return res
path += c
cur = cur.children[c]
# DFS
# since target may be not complete sentence, so we need to do dfs to get all possible sentence
self.dfs(cur,path,res)
return [x[1] for x in sorted(res,key = lambda x: (-x[0],x[1]))[:3]]
def dfs(self,node,path,res):
if node.isWord:
res.append((node.hot,path))
for c in node.children:
self.dfs(node.children[c],path+c,res)
class AutocompleteSystem:
"""
input() tc : O(len(sentence_term) + MlgM) M: matched results sc O(M)
init() tc O(NM) : N: number of sentence M: max length of single sentence sc O(MN)
"""
def __init__(self, sentences: List[str], times: List[int]):
self.root = Trie()
p = self.root
self.search_term = ""
# store historical data
for i,sentence in enumerate(sentences):
p.insert(sentence,times[i])
def input(self, c: str) -> List[str]:
p = self.root
if c != '#':
self.search_term += c
return p.search(self.search_term)
else:
# once a searched sentence complete, store it into trie, reset search term
p.insert(self.search_term,1)
self.search_term = ""
return []
# Your AutocompleteSystem object will be instantiated and called as such:
# obj = AutocompleteSystem(sentences, times)
# param_1 = obj.input(c)