LC 0642 [H] Design Search Autocomplete System - ALawliet/algorithms GitHub Wiki

class TrieNode(object):
    def __init__(self):
        self.children = {}
        self.isEnd = False
        self.data = None
        self.rank = 0
        
class AutocompleteSystem(object):
    def __init__(self, sentences, times):
        self.root = TrieNode()
        self.keyword = ""
        for i, sentence in enumerate(sentences):
            self._addRecord(sentence, times[i])

    def _addRecord(self, sentence, hot):
        cur = self.root
        for c in sentence:
            if c not in cur.children:
                cur.children[c] = TrieNode()
            cur = cur.children[c]
        cur.isEnd = True
        cur.data = sentence
        cur.rank -= hot
    
    def _dfs(self, root):
        res = []
        if root:
            if root.isEnd:
                res.append( (root.rank, root.data) )
            for child in root.children:
                res += self._dfs(root.children[child])
        return res
        
    def _search(self, sentence):
        cur = self.root
        for c in sentence:
            if c not in cur.children:
                return []
            cur = cur.children[c]
        return self._dfs(cur)
    
    def input(self, c):
        res = []
        if c != "#":
            self.keyword += c
            res = self._search(self.keyword)
        else:
            self._addRecord(self.keyword, 1)
            self.keyword = ""
        return [item[1] for item in sorted(res)[:3]] # sorted by hot rank