Train a Next-Word Predictor from Tokenized Sentences

Problem You are given a set of tokenized training sentences (each sentence is a list of words). Train a “next-word predictor”. After training, for any query word q, return the most likely word that immediately follows q in the training data. If q never appears in the training data, or never appears as the first word of any observed bigram, return an empty string (or None depending on the spec; this problem uses empty string). If multiple candidates are tied for the highest frequency, return the lexicographically smallest candidate. Example training data trainingData = [["I", "am", "sam"], ["am", "sam"]] Observed bigrams: ("I" -> "am") ("am" -> "sam") ("am" -> "sam") # from the second sentence So: query = "I" => output "am" query = "am" => output "sam" I/O Format (for a coding prompt) Input Line 1: integer S, number of sentences. Next S lines: one sentence per line, words separated by spaces. Next line: integer Q, number of queries. Next Q lines: one query word q per line. Output For each query, print the predicted next word on its own line; print an empty line if not found. Constraints 1 <= S <= 2e5 Total number of tokens across all sentences T <= 2e6 1 <= Q <= 2e5 Word length <= 50, no spaces. Follow-up discussion Must the space complexity be O(T)? Can you optimize further? What other data structures could be used (e.g., trie, compressed storage, on-disk indexing)? Under a fixed memory budget, roughly how many words/bigrams can be stored? If you want to autocomplete an entire sentence by repeatedly predicting the next token, how do you define stopping conditions, avoid cycles, and handle low-confidence predictions? Example Input 2 I am sam am sam 2 I am Output am sam