Build a bigram next-word predictor that samples the next word in proportion to how often it followed.

Count each word's successors, then sample in proportion to the counts.

import random
from collections import Counter, defaultdict

def train(words):
    counts = defaultdict(Counter)
    for a, b in zip(words, words[1:]):
        counts[a][b] += 1
    return counts

def next_word(counts, w):
    nxt = counts.get(w)
    if not nxt:
        return None              # unseen word
    ws, cs = zip(*nxt.items())
    return random.choices(ws, weights=cs)[0]

Interviewers check the unseen-word case and whether you can sample without random.choices: a cumulative sum plus binary search.