Implement K-Means Clustering From Scratch
Company: Meta
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: medium
Interview Round: Technical Screen
Implement the k-means clustering algorithm from scratch. Given `n` data points in `d` dimensions and a number of clusters `k`, return a cluster label for every point and the final centroids.
The interviewer provides no test cases, so you have to convince yourself and the interviewer that the code is correct by reasoning and with small examples you construct yourself. Unless told otherwise, assume Python with NumPy available but no library implementation of clustering.
```hint Two steps per iteration
Each iteration alternates between two sub-problems, and each has a closed-form answer when the other is held fixed. Writing them as two separate functions makes both easier to check.
```
```hint Test without a test suite
Build one input where the right clusters are obvious by construction, and another that forces the unusual paths in your loop.
```
### Clarifying Questions
- How should centroids be initialized: random data points, k-means++, or centroids passed in as input?
- What stopping rule is expected: a fixed number of iterations, no change in assignments, or centroid movement below a tolerance?
- What should happen if a cluster loses all of its points during an iteration?
- Is a vectorized implementation expected, and are there memory limits on `n`, `k` and `d`?
- Should results be reproducible, for example through a seed argument?
### What a Strong Answer Covers
- Correct assignment and update steps, with vectorized distance computation instead of per-point Python loops in the hot path.
- A stated initialization and stopping rule, with k-means++ or multiple restarts discussed.
- Handling of empty clusters, duplicate points, `k` larger than the number of distinct points, and invalid input.
- Time and memory complexity per iteration, and self-constructed checks in place of the missing test cases.
### Follow-up Questions
- Why does k-means always converge, and why only to a local optimum?
- How would you choose `k` when it is not given?
- How would you scale the algorithm to data that does not fit in memory?
- On what kinds of data does k-means perform poorly, and what would you use instead?
Overview: An ML coding question that asks for a from-scratch implementation of k-means clustering with no provided test cases. It tests the assignment and update steps, vectorized distance computation, initialization, stopping rules, empty-cluster handling and the ability to verify code with self-constructed checks.