ML coding · Probabilistic models
The Viterbi algorithm for a hidden Markov model
A hidden Markov model has S states with initial probabilities start (length S),
transition matrix trans (S × S, row i gives the probabilities of moving from i) and
emission matrix emit (S × V, row i gives the probabilities of each observation from state i).
Given a sequence of observation indices obs, return [path, logp]: the most likely
sequence of hidden states and the natural log of its joint probability with the observations.
All probabilities are positive, and the best path is unique in the tests.
Examples
Input start = [0.6, 0.4], trans = [[0.7, 0.3], [0.4, 0.6]], emit = [[0.5, 0.4, 0.1], [0.1, 0.3, 0.6]], obs = [0, 1, 2]
Output [[0, 0, 1], -4.1917]
The textbook weather example: the hidden states drift from the first to the second as the observations change.
Input start = [0.5, 0.5], trans = [[0.9, 0.1], [0.1, 0.9]], emit = [[0.9, 0.1], [0.2, 0.8]], obs = [0, 0, 1, 1, 1]
Output [[0, 0, 1, 1, 1], -4.192]
Sticky states: one switch, where the observations change.
Input start = [0.3, 0.7], trans = [[0.5, 0.5], [0.5, 0.5]], emit = [[0.8, 0.2], [0.3, 0.7]], obs = [0]
Output [[0], -1.4271]
One observation: the state with the largest start × emission.
Constraints
- 1 ≤ S ≤ 50, 1 ≤ V ≤ 100, 1 ≤ len(obs) ≤ 10⁴
- All probabilities positive
- logp checked to a relative tolerance of 10⁻⁹
Hints
Hint 1
Keep score[j]: the best log-probability of any path ending in state j after the observations so far.
Hint 2
At each step, cand[i, j] = score[i] + log trans[i, j]. The new score[j] is the column maximum plus log emit[j, obs_t]; remember the argmax.
Hint 3
At the end, start from the best final state and follow the stored argmaxes backwards.
Where candidates go wrong
- Multiplying probabilities: with a few hundred observations every path has probability zero in float64.
- Taking the maximum over the wrong axis of cand, so that you choose where to go rather than where you came from.
- Building the path forwards from the argmax at each step, which gives the most likely state at each time, not the most likely sequence.
- An off-by-one between the back-pointer list and the observations.
The idea
The best path ending in state j at time t must consist of the best path ending in some state i at time t − 1, followed by the move i → j. So one vector of best scores per time step is enough, updated with an outer sum and a column maximum: O(S²) per observation instead of enumerating ST paths. Storing the argmax of each column lets you walk back from the best final state.
Logs turn the products into sums and keep every score finite; the argmax is unaffected.
Time O(S²T) · Space O(ST) for the back-pointers
Reference solution
import numpy as np
def viterbi(start, trans, emit, obs):
log_trans = np.log(np.asarray(trans, dtype=float))
log_emit = np.log(np.asarray(emit, dtype=float))
score = np.log(np.asarray(start, dtype=float)) + log_emit[:, obs[0]] # best log-prob ending in each state
back = []
for o in obs[1:]:
cand = score[:, None] + log_trans # cand[i, j]: best path to i, then i -> j
back.append(cand.argmax(axis=0)) # where the best path into j came from
score = cand.max(axis=0) + log_emit[:, o]
path = [int(score.argmax())]
for b in reversed(back):
path.append(int(b[path[-1]]))
return [path[::-1], float(score.max())]Follow-ups an interviewer asks
- The forward algorithm: replace max by log-sum-exp to get the total probability of the observations.
- Posterior decoding: the most likely state at each time, from forward–backward. When does it differ from Viterbi?
- Reduce memory to O(S√T) with checkpointing, or to O(S) if only logp is needed.
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
Viterbi is dynamic programming wearing a probabilistic hat, which makes it a favourite with interviewers who want one problem to test both. It appears in speech, genomics and regime-switching models in finance, and the log-space version is the only one that works on real sequence lengths.
The theory behind it is in the lesson Naive Bayes, EM, Viterbi and Markov chains from scratch.
Keep going
Gaussian naive Bayes classifier
Fit class priors, means and variances, then predict by the largest log posterior, vectorised over test points and classes.
Open the problem → Medium20 minOne EM step for a Gaussian mixture
Compute responsibilities in log space, then refit weights, means and variances: one expectation-maximisation step for a 1-D Gaussian mixture, from scratch..
Open the problem → Medium25 minThe Viterbi algorithm for a hidden Markov model
Find the most likely sequence of hidden states in an HMM with dynamic programming in log space, with back-pointers to recover the path.
Open the problem → Easy12 minStationary distribution of a Markov chain
Solve πP = π with Σπ = 1 by replacing one redundant equation, in NumPy.
Open the problem →