Katkat
Back

Build next-word predictor with O(1) lookup

GoogleData Scientist · Online Assessment
EasyData ScientistOnline AssessmentCoding & Algorithms

Last updated

This problem tests your ability to design a memory-efficient data structure for storing n-gram probabilities and performing fast lookups for the most likely next word.

Problem Statement

You are building the core engine for an autocomplete feature. Given a large corpus of text, you need to store the frequency of word pairs (n-grams) and provide a function that returns the most probable next word for any given input word in constant time O(1).

Requirements

Implement a class NextWordPredictor with the following methods:

  1. train(text: str): Processes a string of text, counts occurrences of adjacent words, and builds the internal data structure.
  2. predict(word: str) -> str | None: Returns the word that most frequently follows the input word. If the word has never been seen or has no followers, return None.

Constraints

  • The corpus size can be large, but the number of unique words (vocabulary size) V is manageable.
  • predict() must perform in O(1) time.
  • If multiple words have the same frequency, return any one of them.

Example

predictor = NextWordPredictor()
predictor.train("the cat chased the mouse and the cat ate the cheese")
# Following 'the': 'cat' (1), 'mouse' (1), 'cheese' (1). 
# Any of these is acceptable.
predictor.predict("the") # Returns 'cat', 'mouse', or 'cheese'
predictor.predict("cat") # Returns 'chased' or 'ate'

Edge Cases

  • Empty string input.
  • Input words that do not exist in the training set.
  • Single-word or very short training inputs.

Hints

Hint 1

Use a hash map (dictionary) where the key is the current word and the value is another hash map of follower counts.

Hint 2

To achieve O(1) lookup, pre-calculate the 'winner' for each word during the training phase or update it dynamically as counts change.

Hint 3

Consider using a secondary hash map to store the 'current best' next word for every entry to avoid re-scanning the follower counts during the predict call.

Follow-up questions

  • How would you modify your design to handle a tie-breaking scenario where you prefer the most recently seen word?
  • If the vocabulary grows to millions of unique words, how would you optimize the memory footprint of your data structure?
  • How would you adapt this model to return the top K most likely words instead of just the top one?