Predict the Most Frequent Next Word in Constant Time

Problem Statement

Implement a class WordPredictor that models word transitions from a collection of tokenized text sequences and answers queries for the most likely next word.

Two words form a transition (w_i, w_{i+1}) if and only if they appear consecutively within the same sequence. Transitions do not cross sequence boundaries (i.e., the last word of sequence k does not transition to the first word of sequence k + 1).

You need to implement the following methods:

• __init__(sequences: list[list[str]]): Initializes the predictor with a list of sequences, where each sequence is a list of lowercase word tokens.
• predict(word: str) -> str: Returns the word that most frequently follows word across all training data. If multiple successor words share the highest frequency, return the lexicographically smallest word. If word has never been followed by another word (or has never been observed at all), return an empty string "".

Queries to predict(word) must run in average O(1) time.

Interface

class WordPredictor:
    def __init__(self, sequences: list[list[str]]):
        pass

    def predict(self, word: str) -> str:
        pass

Examples

Example 1

sequences = [
    ["the", "quick", "brown", "fox"],
    ["the", "lazy", "dog"],
    ["the", "quick", "rabbit"]
]

predictor = WordPredictor(sequences)
predictor.predict("the")    # Output: "quick"
predictor.predict("quick")  # Output: "brown"
predictor.predict("dog")    # Output: ""
predictor.predict("cat")    # Output: ""

Explanation:

• "the" is followed by "quick" twice and "lazy" once. The most frequent successor is "quick".
• "quick" is followed by "brown" once and "rabbit" once. Since counts are tied (1 each), "brown" is returned because "brown" < "rabbit" lexicographically.
• "dog" only appears at the end of a sequence, so it has no successor ("").
• "cat" was never observed, returning "".

Constraints

• 1 <= len(sequences) <= 10^4
• 0 <= len(sequence) <= 10^3
• Total words across all sequences <= 2 * 10^5
• Each word consists only of lowercase English letters ('a' - 'z').
• Up to 10^5 calls will be made to predict.

Edge Cases

• Empty sequences (e.g., []) or single-word sequences (e.g., ["hello"]) generate zero transitions.
• Words that only appear as the last element of sequences have no valid successor.
• Repeated identical transitions within or across sequences must correctly increment transition frequencies.