ML coding · Statistical inference
An exact permutation test for a difference in means
Two groups of integers, a and b. Under the null hypothesis the group
labels are exchangeable. Consider every way of choosing len(a) of the pooled values to form a new first
group, the rest forming the second. Return the fraction of these relabellings whose absolute difference in
group means is at least the observed absolute difference |mean(a) − mean(b)|. The observed labelling is one
of them.
Examples
Input a = [7, 9, 12], b = [3, 4, 6]
Output 0.1
20 relabellings; only the observed split and its mirror image are this extreme, so p = 2/20 = 0.1.
Input a = [1, 2], b = [1, 2]
Output 1
No difference observed: every relabelling is at least as extreme, so p = 1.
Input a = [5], b = [1, 2, 3, 4]
Output 0.4
Five possible splits. 5 alone is the observed one; 1 alone is its mirror image, equally far apart. p = 2/5.
Constraints
- 1 ≤ len(a), len(b)
- len(a) + len(b) ≤ 16
- Values are integers with |v| ≤ 10⁴
- Checked to an absolute tolerance of 10⁻¹²
Hints
Hint 1
Pool the values. A relabelling is a choice of len(a) indices for the first group: itertools.combinations.
Hint 2
The difference in means depends only on the sum of the first group, since the total is fixed.
Hint 3
Compare in integers: multiply the difference by len(a) · len(b) so ties are exact. Count the relabellings at least as extreme and divide by how many there are.
Where candidates go wrong
- Comparing floating-point mean differences with ≥. Two relabellings with the same true difference can round differently, so one counts and the other does not.
- Leaving the observed labelling out of the count. It belongs in the reference set, which is why p can never be 0.
- Counting only one tail when the question asks whether the means differ.
The idea
Let S be the total of the pooled values and s the sum of a candidate first group of size na. Its difference in means is s/na − (S − s)/nb, and multiplying by nanb gives the integer s·nb − (S − s)·na. Compare absolute values of that integer with the observed one, and every tie is counted exactly.
Enumerating C(n, na) subsets is fine up to about 20 values. Past that, draw random relabellings and report (hits + 1) / (draws + 1), which stays a valid p-value.
Time O(C(n, nₐ) · nₐ) · Space O(n)
Reference solution
from itertools import combinations
def permutation_p_value(a, b):
pooled = list(a) + list(b)
na, nb = len(a), len(b)
total = sum(pooled)
# difference in means times na * nb, as an integer, so equal differences compare equal
observed = abs(sum(a) * nb - sum(b) * na)
hits = count = 0
for idx in combinations(range(na + nb), na):
s = sum(pooled[i] for i in idx)
count += 1
if abs(s * nb - (total - s) * na) >= observed:
hits += 1
return hits / countFollow-ups an interviewer asks
- Too many relabellings to enumerate: a Monte Carlo version, and why the estimate is (hits + 1)/(draws + 1).
- A paired design (before and after on the same subjects): flip signs of the differences instead of relabelling.
- Count the relabellings by sum with a dynamic programme, so the exact test runs on larger samples.
- Use the permutation test on the difference in medians, where no t-test applies.
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
It tests whether you understand what a p-value is, rather than which function returns one. The follow-ups (too many relabellings, ties, a one-sided version) separate the candidates who do.
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 →