642. Design Search Autocomplete System

I attempted to create a separate Trie class. It would then become awkward to access the root node of Trie but I was not able to avoid doing so.

Time complexity: assuming inserting N sentences, and the length of each sentence is M

initialization: O(M * N)
input: for #, O(M); otherwise O(27 ^ M) since c is a lowercase English letter or space (can pre-compute the top search results then store it along with each node)

class TrieNode:
    def __init__(self):
        self.children = collections.defaultdict(TrieNode)
        self.sentence = None
        self.freq = 0
        
class AutocompleteSystem:

    def __init__(self, sentences: List[str], times: List[int]):
        self.root = TrieNode()
        self.term = ''
        for i, sentence in enumerate(sentences):
            self.__insert(sentence, times[i])
            
    def input(self, c: str) -> List[str]:
        res = []
        if c != '#':
            self.term += c
            res = self.__search(self.term)
        else:
            self.__insert(self.term, 1)
            self.term = ''
        return [s for t, s in sorted(res)[:3]]
        
    def __insert(self, sentence, times):
        cur = self.root
        for ch in sentence:
            cur = cur.children[ch]
        cur.sentence = sentence
        cur.freq -= times
    
    def __search(self, term):
        cur = self.root
        for ch in term:
            if ch not in cur.children:
                return []
            cur = cur.children[ch]
        return self.__dfs(cur)
    
    def __dfs(self, node):
        res = []
        if node:
            if node.sentence:
                res.append((node.freq, node.sentence))
            for child in node.children:
                res.extend(self.__dfs(node.children[child]))
        return res

Option 2 (better): get top k while inserting

Time complexity: assuming inserting N sentences, and the length of each sentence is M

initialization: O(M * N)
input: O(M)

class TrieNode:
    def __init__(self):
        self.children = collections.defaultdict(TrieNode)
        self.sentence = None
        self.freq = 0
        self.topK = {} # sentence to freq
        
class AutocompleteSystem:

    def __init__(self, sentences: List[str], times: List[int]):
        self.root = TrieNode()
        self.term = ''
        for i, sentence in enumerate(sentences):
            self.__update(sentence, times[i], self.root, 0)
    
    def __update(self, sentence, times, node, idx):        
        if idx == len(sentence):
            node.sentence = sentence
            node.freq -= times
            node.topK[sentence] = node.freq
        else:
            ch = sentence[idx]
            child = node.children[ch]
            candidates = self.__update(sentence, times, child, idx+1)
            for s, t in candidates.items():
                node.topK[s] = t
            
        res = sorted([(t, s) for s, t in node.topK.items()])[:3]
        node.topK = {s:t for t, s in res}
        return node.topK
        
    def __search(self, term):
        cur = self.root
        for ch in term:
            if ch not in cur.children:
                return []
            cur = cur.children[ch]
        return [s for s, t in cur.topK.items()]
            
    def input(self, c: str) -> List[str]:
        if c != '#':
            self.term += c
            return self.__search(self.term)
        else:
            self.__update(self.term, 1, self.root, 0)
            self.term = ''
            return []

642. Design Search Autocomplete System - cocoder39/coco39_LC GitHub Wiki