Q7 - Build a Next-Word Predictor

Scenario: Google predicts your next word constantly. Gboard offers three words above the keyboard before you have typed a letter of them, Gmail finishes your sentence in grey, and Search fills in the rest of your query as you type. Underneath the modern models, the first thing any of these systems learns is embarrassingly simple: which word tends to come after which other word.

You are handed a corpus that has already been tokenized into sentences, and asked to build that model. The version you are building looks only at pairs of words that sit next to each other, counts how often each word follows each other word, and answers questions about it.

The input, precisely

• sentences is a list of sentences. Each sentence is a list of word tokens, already split for you, in the order they were written. Nothing needs cleaning: no punctuation to strip, no case to fold. Two tokens are the same word only if they are exactly equal as strings, so The and the are different words.
• queries is a list of word tokens to ask about. The same word may appear in it more than once, and a query word need not appear in the corpus at all.
• seed is a single non-negative integer, and it is only used in Part 3.

The output, precisely

Return a list with exactly one answer per query, in the same order as queries. Each answer is either a word token or None, meaning "no prediction."

The rules that decide every answer

• A pair only counts inside one sentence. If a sentence ends with ham and the next sentence begins with I, that is not a pair. Sentences never run into each other.
• The last token of a sentence has no follower. It contributes nothing as a source word from that position.
• Counts pool across the whole corpus. If I am appears in three different sentences, am has a count of 3 as a follower of I.
• Ties go to the lexicographically smallest word, compared as plain strings.
• The answer is None in exactly two situations, and it is worth saying both out loud: the word never appears in the corpus at all, or it appears only ever as the last token of a sentence, so it was never observed with anything after it. Both mean "I have never seen what follows this word," and both return None.

Constraints

0 <= number of sentences <= 100000
0 <= total tokens across all sentences <= 200000
0 <= number of queries <= 200000
0 <= seed < 2147483648
every token is a non-empty string

An empty corpus, an empty query list, and sentences of length 0 or 1 are all legal inputs. A corpus of nothing but one-word sentences produces no pairs at all, so every query answers None.

Python        None           (typed as Optional[str])
JavaScript    null
TypeScript    null           (typed as string | null)
Java          null           (inside a String[])
C#            null           (inside a string?[])
C++           std::nullopt   (std::optional<std::string>)
Go            ""             (tokens are never empty, so this is safe)
Kotlin        null           (typed as String?)

Part 1: The Most Likely Next Word

For each query word, return the word that most often follows it. One function: read the corpus, then answer every query.

Input:    sentences = [["I", "am", "Sam"],
                      ["Sam", "I", "am"],
                      ["I", "like", "green", "eggs", "and", "ham"]]
          queries   = ["I", "Sam", "ham", "green", "missing"]

Output:   ["am", "I", None, "eggs", None]

Why:      "I"        followed by "am" twice, "like" once  ->  "am"
          "Sam"      followed by "I" once, nothing else    ->  "I"
          "ham"      appears only as a last token          ->  None
          "green"    followed by "eggs" once               ->  "eggs"
          "missing"  never appears in the corpus           ->  None

Two details in that example are the ones interviewers actually probe. Sam is the last token of sentence 1 and the first token of sentence 2, and only the second of those produces a pair, because pairs never cross a sentence boundary. And ham shows that "appears in the corpus" and "has a follower" are different questions.

Input:    sentences = [["a", "b"], ["a", "c"]]
          queries   = ["a"]

Output:   ["b"]

Why:      "b" and "c" each follow "a" exactly once, so the counts tie.
          The tie goes to the lexicographically smaller word, so "b" wins.

Step 1: Slide a two-word window over each sentence

The model is built from pairs, so the first job is to produce them. Walk each sentence with a window of two and emit (word, follower) for every adjacent position. A sentence of length L yields L - 1 pairs, which is 0 for a sentence of one word and 0 for an empty sentence.

Here is Example 1's corpus turned into pairs:

sentence 1   ["I", "am", "Sam"]
               I -> am        am -> Sam

sentence 2   ["Sam", "I", "am"]
               Sam -> I       I -> am

sentence 3   ["I", "like", "green", "eggs", "and", "ham"]
               I -> like      like -> green    green -> eggs
               eggs -> and    and -> ham

Nine pairs from eleven tokens, because each of the three sentences loses its last token as a source. The words Sam and ham are the ones that ended a sentence, and you can already see that ham never shows up on the left of an arrow.

Step 2: Tally the pairs into a map of maps

Now count. The natural home for a pair tally is a map keyed by the source word, whose value is another map from follower to how many times that follower was seen. People call this a bigram table.

followerCounts

   "I"      ->  { "am": 2, "like": 1 }
   "am"     ->  { "Sam": 1 }
   "Sam"    ->  { "I": 1 }
   "like"   ->  { "green": 1 }
   "green"  ->  { "eggs": 1 }
   "eggs"   ->  { "and": 1 }
   "and"    ->  { "ham": 1 }

Notice what this structure gets right for free. ham has no key at all, which is the same state as a word that never appeared, and both are exactly the None case. You do not need a separate check for "appeared but only at the end of a sentence," because a word that only ever ended a sentence was never used as a source and so was never inserted.

Step 3: Pick the winner for one query word

Look the query up. If there is no entry, answer None. Otherwise scan that word's followers once, keeping the best you have seen under the rule "higher count wins, and on equal counts the lexicographically smaller word wins."

query "I"   ->  { "am": 2, "like": 1 }

   candidate "am"    count 2   best is empty  ->  take it   best = ("am", 2)
   candidate "like"  count 1   1 beats 2?  no
                               1 ties 2?   no  ->  keep      best = ("am", 2)

   answer: "am"

One comparison expresses both halves of that rule at once, and it is worth writing it down before you code it, because getting it backwards is the single most common bug in Part 1:

# take the candidate if it wins on count, or ties on count and wins on spelling
if count > best_count or (count == best_count and follower < best_follower):
    best_follower, best_count = follower, count

Seeding best_follower with the empty string and best_count with 0 makes the first candidate always win, since every real count is at least 1. That removes the "is this the first one?" branch entirely.

from typing import Dict, List, Optional


def predict_next_words(
    sentences: List[List[str]], queries: List[str]
) -> List[Optional[str]]:
    follower_counts: Dict[str, Dict[str, int]] = {}
    for sentence in sentences:
        for index in range(len(sentence) - 1):
            word, follower = sentence[index], sentence[index + 1]
            counts = follower_counts.setdefault(word, {})
            counts[follower] = counts.get(follower, 0) + 1

    answers: List[Optional[str]] = []
    for query in queries:
        counts = follower_counts.get(query)
        if not counts:
            answers.append(None)
            continue
        best_follower, best_count = "", 0
        for follower, count in counts.items():
            if count > best_count or (
                count == best_count and follower < best_follower
            ):
                best_follower, best_count = follower, count
        answers.append(best_follower)
    return answers

One line in the C++ version is worth stealing even if C++ is not your language. Writing the loop as index + 1 < sentence.size() instead of index < sentence.size() - 1 is what keeps an empty sentence from blowing up, because size() is unsigned and 0 - 1 wraps to a gigantic number. In the languages where lengths are signed, 0 - 1 is just -1 and the loop never runs, so both spellings are safe. The constraints explicitly allow an empty sentence, so this is a real input, not a hypothetical.

Part 2: Answering in Constant Time

Same answers, different cost. Redesign the preprocessing so that answering one query is O(1) on average, a hash lookup and nothing more.

Before you touch the keyboard, put a number on why this is being asked. Part 1 rescans a word's whole follower set on every query. Let k be how many distinct words have been seen after the query word. At the stated limits, 200,000 tokens can put roughly 100,000 distinct followers behind a single popular word, and 200,000 queries can all ask about that same word. That is about 2 * 10^10 comparisons in the query phase, for a model that only took 200,000 steps to build. The work is in completely the wrong place.

Input:    sentences = [["a", "c"], ["a", "b"], ["a", "b"],
                       ["a", "c"], ["c", "d"]]
          queries   = ["a", "c"]

Output:   ["b", "d"]

Why:      "a" is followed by "b" twice and "c" twice, so the counts tie at 2,
          and "b" wins the tie on spelling.
          "c" is followed by "d" once. It is also the last token of two
          sentences, and those two positions contribute nothing.

Step 1: Move the winner out of the query and into the training loop

The fix is not a faster scan, it is no scan. The answer to "what follows a" cannot change between queries, because training is finished before any query arrives. So decide it once, while you are counting, and store it.

Keep two maps next to the tally:

bestFollower   word -> the word that currently wins for it
bestCount      word -> how many times that winner was seen

A query is then one lookup in bestFollower. A missing key is the None answer, exactly as before.

Step 2: Update the winner as each pair arrives

This works because of two facts about the counting loop, and stating them is most of what the interviewer wants to hear. Counts only ever go up, and each pair changes exactly one count. So when you bump follower to newCount, only three things can be true:

• newCount beats the stored count. follower is now strictly the highest, so it is the new winner.
• newCount ties the stored count. There are two words at the top, and the smaller spelling takes it.
• newCount is below the stored count. Nothing else moved, so the stored winner is still the winner.

That is the same comparison from Part 1, applied once per pair instead of once per query. And you never have to handle "what if the word I just bumped is already the winner," because if it were the winner its stored count was newCount - 1, so the first branch fires and it simply keeps the title with a higher number.

Walk Example 3's five pairs in order and every branch shows up:

   pair     counts so far    best before   what happens            best after
1  a -> c   a: {c:1}         ("", 0)       1 beats 0               ("c", 1)
2  a -> b   a: {c:1, b:1}    ("c", 1)      ties 1, "b" < "c"       ("b", 1)
3  a -> b   a: {c:1, b:2}    ("b", 1)      2 beats 1               ("b", 2)
4  a -> c   a: {c:2, b:2}    ("b", 2)      ties 2, "c" loses       ("b", 2)
5  c -> d   c: {d:1}         ("", 0)       1 beats 0               ("d", 1)

   bestFollower = { "a": "b", "c": "d" }      queries ["a", "c"]  ->  ["b", "d"]

Step 4 is the one to point at out loud. c catches up to b on count, and the answer still does not move, because a tie is settled by spelling and b was already ahead there. A candidate who writes >= instead of > in the first half of the comparison gets c here and fails this exact example.

# runs once per pair, inside the counting loop
current_follower = self.best_follower.get(word, "")
current_count = self.best_count.get(word, 0)
if new_count > current_count or (
    new_count == current_count and follower < current_follower
):
    self.best_follower[word] = follower
    self.best_count[word] = new_count

Note what "O(1) average" is actually promising. It is the hash table's average-case lookup, not a worst-case guarantee: a pathological set of keys that all collide degrades it. Saying that unprompted costs you one sentence and is the difference between a candidate who repeats the phrase and one who knows what it means.

from typing import Dict, List, Optional


class NextWordModel:
    def __init__(self, sentences: List[List[str]]):
        follower_counts: Dict[str, Dict[str, int]] = {}
        self.best_follower: Dict[str, str] = {}
        self.best_count: Dict[str, int] = {}
        for sentence in sentences:
            for index in range(len(sentence) - 1):
                word, follower = sentence[index], sentence[index + 1]
                counts = follower_counts.setdefault(word, {})
                new_count = counts.get(follower, 0) + 1
                counts[follower] = new_count
                current_follower = self.best_follower.get(word, "")
                current_count = self.best_count.get(word, 0)
                if new_count > current_count or (
                    new_count == current_count and follower < current_follower
                ):
                    self.best_follower[word] = follower
                    self.best_count[word] = new_count

    def predict(self, word: str) -> Optional[str]:
        return self.best_follower.get(word)

The followerCounts tally is still there, and it has to be: the winner update reads the count it just produced. What changed is that nothing reads the tally after training ends. If memory were tight you could even drop it once the model is built, since bestFollower is all a query needs.

Part 3: Sampling by Frequency

Always returning the top word makes a very dull keyboard. A suggestion system that samples from the distribution it learned feels alive, and it surfaces the less common words too, just less often. So instead of the winner, return a random next word, drawn in proportion to how many times each follower was actually observed.

Random and testable do not usually go together, so the generator is pinned down exactly. Follow these five rules to the letter and your output is reproducible:

1. Keep an integer state, set to seed before the first query.
2. For each query word that has at least one follower, turn the generator once: state = (1103515245 * state + 12345) % 2147483648.
3. Let target = state % total_count(word), where total_count(word) is the total number of followers observed for that word, counting repeats. For { "am": 2, "like": 1 } that is 3, not 2.
4. Sort that word's candidate followers lexicographically and build cumulative counts in that order. Return the first word whose cumulative count is strictly greater than target.
5. If a query word has no follower, return None and do not turn the generator.

Input:    sentences = [["I", "am", "Sam"],
                      ["Sam", "I", "am"],
                      ["I", "like", "green", "eggs", "and", "ham"]]
          queries   = ["I", "I", "Sam", "ham"]
          seed      = 1

Output:   ["am", "am", "I", None]

Why:      "I" has the distribution { "am": 2, "like": 1 }, so "am" holds 2 of
          its 3 slots. The generator lands on an "am" slot both times.
          "Sam" has only one follower, so every draw returns "I".
          "ham" has no follower: None, and the generator does not turn.

Input:    sentences = [["a", "c"], ["a", "b"]]
          queries   = ["a", "a", "a"]
          seed      = 0

Output:   ["c", "b", "c"]

Why:      "a" is followed once by "b" and once by "c", an even split over
          2 slots. Three draws from the same seeded sequence give "c",
          then "b", then "c". Note that the top word by frequency is "b"
          (it wins Part 2's tie), yet the first draw here returns "c".
          Sampling ignores the winner.

Step 1: Lay each word's followers out as a ruler of slots

Sampling by frequency is easiest to see as a ruler. Give every observation its own slot, so a follower seen twice takes two slots, then pick a slot uniformly at random. The word that owns the chosen slot is the answer, and it is chosen in exact proportion to its count, which is what "empirical frequency" means.

You do not build that ruler literally, because a word seen 50,000 times would cost 50,000 entries. You store the cumulative counts in lexicographic order instead, which is the same ruler described by its boundaries:

word "I"     observed followers: { "am": 2, "like": 1 }

sorted lexicographically       am      like
count                           2         1
cumulative                      2         3       total = last cumulative = 3

the ruler this describes
   slot      0       1       2
   owner    am      am      like

target = state % 3 picks a slot in 0, 1, 2. Slots 0 and 1 belong to am, slot 2 belongs to like, which is exactly the 2-to-1 split the corpus showed.

Two things about rule 4 are easy to get wrong and both are visible here. The order is lexicographic, not by count, so am before like because of spelling, and it would still be am first even if like had been seen ten times. And the comparison is strictly greater, not greater-or-equal: at target = 1, am's cumulative 2 is greater than 1 so am wins, and using >= would have handed slot 2 to am as well and made like unreachable.

Step 2: Turn the generator, once per answerable query

That recurrence is a linear congruential generator, the same shape the old C rand() used. Nothing about it needs to be good, it only needs to be identical on your machine and the grader's.

Here is Example 4 turned one step at a time. Read the state column as "after turning the generator for this query":

   query    state after turn   total   target   answer   why
1  "I"      1103527590            3       0      am       cumulative 2 beats 0
2  "I"      377401575             3       0      am       cumulative 2 beats 0
3  "Sam"    662824084             1       0      I        only one slot exists
4  "ham"    not turned            -       -      None     no follower at all

And Example 5, where the same word gives three different draws because the state keeps moving:

ruler for "a":   slot 0 = b,   slot 1 = c

   query    state        total   target   cumulative b = 1   answer
1  "a"      12345          2       1      not above 1        c
2  "a"      1406932606     2       0      above 0            b
3  "a"      654583775      2       1      not above 1        c

Step 3: The two things that actually break this

Do not turn the generator for a None answer. Rule 5 is not a detail, it is a fork in the output. Every draw after a skipped word reads a different state, so one extra turn early on corrupts every answer that follows it. On sentences = [["a", "c"], ["a", "b"]] with queries = ["zzz", "a"] and seed = 0:

correct    "zzz" has no follower, so the state stays 0.
           Then "a" turns it to 12345, target 1, answer "c".
           ->  [None, "c"]

buggy      "zzz" turns the state to 12345 and throws the draw away.
           Then "a" turns it to 1406932606, target 0, answer "b".
           ->  [None, "b"]

The first answer is None either way, which is what makes this so easy to miss: nothing looks wrong until the query after it.

The multiply overflows, in more languages than you expect. 1103515245 * state with state just under 2^31 reaches about 2.4 * 10^18. That number is fine in a signed 64-bit integer, which tops out near 9.2 * 10^18, and catastrophic in a signed 32-bit one. So long in Java, long long in C++, int64 in Go, Long in Kotlin, and long in C#.

JavaScript and TypeScript have a subtler version of the same problem, and it is worth knowing because nothing throws. A number is a 64-bit float, which holds integers exactly only up to 2^53, about 9.0 * 10^15. The product is a thousand times larger than that, so it gets rounded and the sequence drifts. Starting from seed = 1, the second state should be 377401575, and plain number arithmetic produces 377401600. Reaching for BigInt costs one line and removes the whole class of problem.

from typing import Dict, List, Optional

MULTIPLIER, INCREMENT, MODULUS = 1103515245, 12345, 2147483648


def sample_next_words(
    sentences: List[List[str]], queries: List[str], seed: int
) -> List[Optional[str]]:
    follower_counts: Dict[str, Dict[str, int]] = {}
    for sentence in sentences:
        for index in range(len(sentence) - 1):
            word, follower = sentence[index], sentence[index + 1]
            counts = follower_counts.setdefault(word, {})
            counts[follower] = counts.get(follower, 0) + 1

    sorted_followers: Dict[str, List[str]] = {}
    cumulative_counts: Dict[str, List[int]] = {}
    for word, counts in follower_counts.items():
        followers = sorted(counts)
        running, total = [], 0
        for follower in followers:
            total += counts[follower]
            running.append(total)
        sorted_followers[word] = followers
        cumulative_counts[word] = running

    state = seed
    answers: List[Optional[str]] = []
    for query in queries:
        cumulative = cumulative_counts.get(query)
        if cumulative is None:
            answers.append(None)
            continue
        state = (MULTIPLIER * state + INCREMENT) % MODULUS
        target = state % cumulative[-1]
        for position, reached in enumerate(cumulative):
            if reached > target:
                answers.append(sorted_followers[query][position])
                break
    return answers

One upgrade to mention if you finish with time on the clock. cumulative is sorted ascending by construction, so "the first entry strictly greater than target" is exactly an upper bound, and the linear scan can be a binary search:

# alternative to the scan: O(log k) instead of O(k) per query
from bisect import bisect_right

position = bisect_right(cumulative, target)
answers.append(sorted_followers[query][position])

std::upper_bound in C++ and sort.Search in Go do the same job. It only shows up in the timings when a single word has a very large follower set, so the honest framing is "worth it if the distribution is skewed, and not worth complicating the first version over."

All Three Parts in One Model

Each part on its own re-reads the corpus. Built together, the corpus is read once and the single counting loop feeds two different indexes, one for each kind of prediction:

one pass over the sentences
        |
        +--> bestFollower[word]       the running winner, kept up
        |                             to date as counts rise      ->  predict()
        |
        +--> sortedFollowers[word]    followers in lexicographic
        |                             order                       ->  sample()
        +--> cumulativeCounts[word]   running totals over that
                                      same order

both queries are then a hash lookup, and sample() adds one turn
of the generator

The only ordering constraint is that the cumulative ruler is built after counting finishes, since it needs the final totals. The running winner, by contrast, can be maintained inside the counting loop, because counts only ever rise and each pair changes exactly one of them.

from typing import Dict, List, Optional

MULTIPLIER, INCREMENT, MODULUS = 1103515245, 12345, 2147483648


class NextWordModel:
    def __init__(self, sentences: List[List[str]], seed: int = 0):
        follower_counts: Dict[str, Dict[str, int]] = {}
        self.best_follower: Dict[str, str] = {}
        best_count: Dict[str, int] = {}
        for sentence in sentences:
            for index in range(len(sentence) - 1):
                word, follower = sentence[index], sentence[index + 1]
                counts = follower_counts.setdefault(word, {})
                new_count = counts.get(follower, 0) + 1
                counts[follower] = new_count
                current_count = best_count.get(word, 0)
                if new_count > current_count or (
                    new_count == current_count
                    and follower < self.best_follower.get(word, "")
                ):
                    self.best_follower[word] = follower
                    best_count[word] = new_count

        self.sorted_followers: Dict[str, List[str]] = {}
        self.cumulative_counts: Dict[str, List[int]] = {}
        for word, counts in follower_counts.items():
            followers = sorted(counts)
            running, total = [], 0
            for follower in followers:
                total += counts[follower]
                running.append(total)
            self.sorted_followers[word] = followers
            self.cumulative_counts[word] = running
        self.state = seed

    def predict(self, word: str) -> Optional[str]:
        """The most frequent follower, ties to the lexicographically smallest."""
        return self.best_follower.get(word)

    def sample(self, word: str) -> Optional[str]:
        """A follower drawn in proportion to how often it was observed."""
        cumulative = self.cumulative_counts.get(word)
        if cumulative is None:
            return None
        self.state = (MULTIPLIER * self.state + INCREMENT) % MODULUS
        target = self.state % cumulative[-1]
        for position, reached in enumerate(cumulative):
            if reached > target:
                return self.sorted_followers[word][position]
        return None

Complexity Analysis

Let T be the total number of tokens, Q the number of queries, and k the number of distinct followers of a particular word.

• Part 1. Training is O(T): one pass, one hash write per pair. Each query is O(k) for the scan, so the query phase is O(Q * k) and the whole thing is O(T + Q * k). With one dominant word, k can approach T / 2, which is where the 2 * 10^10 figure came from.
• Part 2. Training is still O(T), since the winner update is a fixed amount of work per pair. Each query is O(1) on average, one hash lookup, so the total is O(T + Q). This is the version that scales, and it is faster than Part 1 by doing strictly less work, not by doing work earlier.
• Part 3. Training is O(T) to count plus the sorting. Summed over every source word that is O(T log T) in the worst case, and much closer to O(T) in real text, where most words have only a handful of distinct followers. Each query is O(k) for the walk along the cumulative array, or O(log k) with the binary search.
• Memory for all three is O(P), where P is the number of distinct pairs in the corpus, and P is at most T - 1. Part 3 holds one extra copy of each word's follower list in sorted order, so it is a constant factor above the others, not a different order.

Final Notes

1. Where you put the work is the whole question. All three parts count the same pairs. What moves is the boundary between training and querying: Part 1 ranks at query time, Part 2 ranks during training, Part 3 precomputes a ruler and then only indexes into it. Interviewers ask it this way because the choice of where to spend the work is the actual engineering decision.
2. Two different words give you None, and the map handles both. Never seen at all, or seen only at the end of a sentence, are the same thing to the model: no key. If you catch yourself adding a second check for "appears but has no follower," step back and look at your data structure, because it already knows.
3. The tie rule is a two-clause comparison, and the second clause is where people lose the example. Higher count wins outright, and only on an exact tie does the smaller spelling decide. Seed the running best with an empty string and a count of zero and the first candidate always wins, so the "is this the first?" branch disappears.
4. The incremental winner works because counts only rise. Each pair changes exactly one count, and it can only go up, so a single comparison per pair keeps the leader current with no rescanning. Be ready to say that sentence, because "why is that correct?" is the natural follow-up and it is the reason the redesign is allowed at all.
5. A skipped draw is a corrupted sequence. Turning the generator on a word that has no follower changes the answer to every query after it, and the None it returned looks perfectly correct. Bugs that only surface downstream are the ones worth writing a test for on the spot.
6. Check the width of your integers before the interviewer does. The multiply reaches about 2.4 * 10^18: fine in 64 bits, ruinous in 32, and silently wrong in a JavaScript number. Naming the range out loud takes five seconds and is the kind of thing that separates a candidate who has shipped numeric code from one who has not.

Predicting the next word is where every language model starts, and this is the first rung of that ladder: count what you have seen, then either take the most likely thing or roll for it. Get the counting and the ranking to stay in separate places, and the version that samples is a twenty-minute extension instead of a rewrite. 🚀