Skip to content
Work Practice Products Coding interviews Machine learning coding Quant developer HFT course Blog Method Why me About Enquire Book a call

Practice · ML coding

k-means, kNN, PCA and decision trees from scratch

Older than deep learning and still asked, because each one fits in twenty lines and each one hides a decision about ties, signs or empty sets.


Pairwise distances without a loop

kNN and k-means both need the distance from every point in one set to every point in another. Broadcasting gives it in one line:

D = ((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=2)    # (len(A), len(B))

A[:, None, :] has shape (n, 1, d) and B[None, :, :] has shape (1, m, d); their difference broadcasts to (n, m, d). That costs O(nmd) memory. When it is too much, expand the square instead: ‖a − b‖² = ‖a‖² − 2a·b + ‖b‖², which needs only the (n, m) result and one matrix product. The expansion can lose precision when the points are far from the origin and close to each other, and can even go slightly negative; clip at zero before any square root.

For ranking neighbours, compare squared distances. The square root changes no ordering and costs time.

k-nearest neighbours

There is no training step. To predict, find the k closest training points and take a majority vote. The work is all in the details an interviewer will probe: what happens when two training points are equally far away, and when two labels get the same number of votes. A good answer states a rule for each before writing code. np.argsort(kind='stable') keeps equal distances in index order, and a dictionary filled in order of distance remembers which label was seen first.

k-means

Lloyd's algorithm alternates two steps until nothing changes: assign each point to its nearest centroid, then move each centroid to the mean of its points. Each step can only lower the total squared distance, so the algorithm stops, but at a local minimum that depends on where it started. That is why real implementations run it several times from k-means++ starts.

The question to answer before coding: what happens to a centroid that loses all its points? Its mean is undefined. Leaving it where it is, reseeding it at the farthest point, or dropping it are all defensible; choose one and say which.

Lloyd's algorithm on ten points with k = 2.

  1. Step 1 / 7
    c1 (1, 2)c2 (2.5, 1.5)

    Ten points and two starting centroids, both picked from the left group: a deliberately poor start.

  2. Step 2 / 7
    c1 (1, 2)c2 (2.5, 1.5)

    Iteration 1, assign: each point joins its nearest centroid. c1 has 4 points, c2 has 6.

  3. Step 3 / 7
    c1 (1.575, 2.3)c2 (6.3833, 5.5)

    Iteration 1, update: each centroid moves to the mean of its points.

  4. Step 4 / 7
    c1 (1.575, 2.3)c2 (6.3833, 5.5)

    Iteration 2, assign: each point joins its nearest centroid. c1 has 5 points, c2 has 5.

  5. Step 5 / 7
    c1 (1.76, 2.14)c2 (7.16, 6.3)

    Iteration 2, update: each centroid moves to the mean of its points.

  6. Step 6 / 7
    c1 (1.76, 2.14)c2 (7.16, 6.3)

    Iteration 3, assign: each point joins its nearest centroid. c1 has 5 points, c2 has 5.

  7. Step 7 / 7
    c1 (1.76, 2.14)c2 (7.16, 6.3)

    Iteration 3, update: no centroid moves. The assignment is stable and the algorithm stops.

PCA

Centre the data, then find the directions of greatest variance. Two routes give the same answer:

  • the eigenvectors of the covariance matrix XᵀX/(n − 1), sorted by eigenvalue;
  • the right singular vectors of the centred X, from np.linalg.svd.

The SVD is the better one: it never forms XᵀX, which squares the condition number. Either way the result has a sign ambiguity: if v is a principal direction, so is −v. Two correct programs can disagree on every sign, so any test, and any production pipeline that compares runs, needs a convention. A common one is to make the largest-magnitude entry of each component positive.

The decision-tree split

A tree is built from one operation, repeated: find the feature and threshold that best separate the labels. Gini impurity, 1 − Σ pc², is zero for a pure node. The quality of a split is the impurity of the two children, weighted by their sizes. Sort each feature once, sweep the threshold from left to right while moving one point at a time from the right child's class counts to the left child's, and every candidate costs O(1). That turns O(n²) per feature into O(n log n).

What interviewers listen for

  • Rules for ties, stated up front.
  • Squared distances when only the order matters.
  • SVD rather than the covariance matrix, and a sign convention.
  • Thresholds placed between distinct values, never on a value shared by two points.

Practise it