ML coding · Classic algorithms
Principal component analysis
Given data X (n × d) and k ≤ min(n, d), return
[components, projected]:
components(k × d): the top k principal directions of the centred data, as unit-length rows, in decreasing order of variance;- each row's sign fixed so that its entry of largest absolute value is positive (the tests avoid ties between entries of opposite sign);
projected(n × k): the centred data expressed in those directions.
The tests guarantee the top k variances are distinct.
Examples
Input X = [[1, 1], [2, 2], [3, 3]], k = 1
Output [[[0.7071, 0.7071]], [[-1.4142], [0], [1.4142]]]
All the variance lies along the diagonal (0.7071, 0.7071).
Input X = [[0, 0], [2, 0], [0, 1], [2, 1]], k = 2
Output [[[1, 0], [0, 1]], [[-1, -0.5], [1, -0.5], [-1, 0.5], [1, 0.5]]]
Variance 4/3 along x and 1/3 along y: the axes, in that order.
Input X = [[1, -2], [-1, 2], [3, -6]], k = 1
Output [[[-0.4472, 0.8944]], [[0], [4.4721], [-4.4721]]]
The direction (0.4472, −0.8944) and its negative are both right. The sign rule makes the larger entry positive.
Constraints
- 2 ≤ n ≤ 1000, 1 ≤ d ≤ 50, 1 ≤ k ≤ min(n, d)
- The top k + 1 singular values are distinct (k + 1 ≤ min(n, d))
- Checked to a relative tolerance of 10⁻⁷
Hints
Hint 1
Subtract the column means first. PCA on uncentred data finds the direction of the mean, not of the variance.
Hint 2
np.linalg.svd(Xc, full_matrices=False) returns Vt; its rows are the principal directions, already sorted.
Hint 3
For the sign: find each row's index of largest |value| with argmax, and multiply the row by the sign of that entry. Then project with Xc @ W.T.
Where candidates go wrong
- Forgetting to centre.
- Taking columns of Vt instead of rows (or rows of V instead of columns).
- Forming the covariance matrix and calling np.linalg.eig: eig does not sort its output and may return complex dtype. If you go that way, use eigh and reverse the order.
- Projecting the uncentred data.
The idea
Write the centred data as Xc = UΣVᵀ. Then XcᵀXc = VΣ²Vᵀ, so the columns of V are the eigenvectors of the covariance matrix, with variances σi²/(n − 1), already sorted. Taking the top k rows of Vᵀ gives the components without ever forming XᵀX, which would square the condition number.
Singular vectors are defined only up to sign. Flipping each component so that its largest entry is positive makes the output unique, and then the projection XcWᵀ is unique too.
Time O(nd · min(n, d)) · Space O(nd)
Reference solution
import numpy as np
def pca(X, k):
X = np.asarray(X, dtype=float)
Xc = X - X.mean(axis=0) # centre each feature
_, _, Vt = np.linalg.svd(Xc, full_matrices=False)
W = Vt[:k].copy() # rows: principal directions, largest variance first
idx = np.abs(W).argmax(axis=1)
W *= np.sign(W[np.arange(k), idx])[:, None] # sign rule: largest entry positive
return [W, Xc @ W.T]Follow-ups an interviewer asks
- Return the fraction of variance explained by each component.
- Reconstruct the data from the k components and report the reconstruction error. How does it relate to the discarded singular values?
- Randomised PCA for n = 10⁷: project onto a random subspace first.
- Whitening: scale each projected coordinate to unit variance. When does that hurt?
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
PCA is the standard test of linear algebra in an ML interview: centring, the link between the SVD and the covariance matrix, and an awareness that eigenvectors have no preferred sign. Quant research interviews ask it constantly, because the first principal component of a set of returns is the market.
The theory behind it is in the lesson k-means, kNN, PCA and decision trees from scratch.
Keep going
k-nearest neighbours classifier
Predict labels by majority vote among the k closest training points, with explicit tie rules.
Open the problem → Medium20 mink-means clustering (Lloyd's algorithm)
Lloyd's algorithm from given starting centroids: vectorised assignment, mean update, empty clusters and convergence, written from scratch in NumPy..
Open the problem → Medium15 minPrincipal component analysis
Centre the data, take the SVD, keep the top k directions with a fixed sign, and project.
Open the problem → Medium25 minBest decision-tree split by Gini impurity
Find the feature and threshold that minimise weighted Gini impurity, with a sorted sweep in O(n log n) per feature.
Open the problem →