ML coding · Statistical inference
The autocorrelation function of a series
Given a series x of length n and an integer max_lag with
0 ≤ max_lag < n, return the sample autocorrelations r0, …, rmax_lag, where
rk = Σt=0n−k−1 (xt − x̄)(xt+k − x̄) / Σt=0n−1 (xt − x̄)²
and x̄ is the mean of the whole series. The series is not constant.
Examples
Input x = [1, 2, 3, 4, 5], max_lag = 2
Output [1, 0.4, -0.1]
A trend: strong positive correlation at lag 1, weaker at lag 2. r₀ is always 1.
Input x = [1, -1, 1, -1, 1, -1], max_lag = 1
Output [1, -0.8333]
Alternating signs: r₁ is strongly negative.
Input x = [0.5, 0.1, -0.3, 0.2], max_lag = 0
Output [1]
Lag 0 only.
Constraints
- 2 ≤ n ≤ 10⁵
- 0 ≤ max_lag < n
- x is not constant
- Checked to a relative tolerance of 10⁻⁹
Hints
Hint 1
Subtract the mean once, from the whole series. Every lag uses the same deviations.
Hint 2
The numerator at lag k is a dot product of the deviations with themselves shifted by k.
Hint 3
d[:n − k] @ d[k:] for each k, divided by d @ d. For long series and many lags, an FFT computes all of them in O(n log n).
Where candidates go wrong
- Using a separate mean for each of the two overlapping segments: that is a different estimator, closer to a Pearson correlation of the segments.
- Dividing each lag by n − k instead of using the same denominator. It reduces the bias but can produce an autocorrelation sequence that is not a valid correlation structure.
- Looping over t in Python for every lag: O(n · max_lag) interpreted steps instead of vectorised dot products.
The idea
Centre the series once: d = x − x̄. The lag-k numerator pairs each deviation with the one k steps later, which is the dot product of d without its last k entries and d without its first k. The denominator is the lag-0 numerator, so r₀ = 1 by construction and every |rk| ≤ 1.
Using one denominator for all lags makes the sequence positive semi-definite, the property any autocorrelation of a stationary process has. Dividing each lag by its own number of terms breaks that, which matters when the estimates feed a model fit.
Time O(n · max_lag), or O(n log n) by FFT · Space O(n)
Reference solution
import numpy as np
def autocorrelation(x, max_lag):
d = np.asarray(x, dtype=float)
d = d - d.mean() # one mean for the whole series
n = len(d)
denom = d @ d # the same denominator at every lag
return [float(d[:n - k] @ d[k:] / denom) for k in range(max_lag + 1)]Follow-ups an interviewer asks
- Compute every lag at once with an FFT: pad to length 2n, transform, take the squared modulus, invert.
- The Ljung–Box statistic: n(n + 2) Σ r_k² / (n − k), compared with a chi-squared distribution.
- The partial autocorrelation function, from the Durbin–Levinson recursion.
- Check the autocorrelation of squared returns for volatility clustering.
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
A warm-up in research interviews: it shows whether you know the standard estimator (one mean, one denominator) and can vectorise a lagged sum without a Python loop over time.
The theory behind it is in the lesson Statistical inference in code: standard errors, permutation tests, the Kalman filter.
Keep going
The autocorrelation function of a series
Compute the sample autocorrelation of a time series at lags 0 to k, vectorised, with the standard estimator.
Open the problem → Medium20 minRegression coefficients with standard errors and t-statistics
Fit least squares with QR, then report each coefficient's standard error and t-statistic.
Open the problem → Medium20 minAn exact permutation test for a difference in means
Compute the exact two-sided p-value for a difference in means by trying every relabelling of the pooled data.
Open the problem → Medium20 minA one-dimensional Kalman filter
Track a drifting level through noisy measurements with the Kalman filter: predict, compute the gain, update.
Open the problem →