Vectorized 1-Nearest-Neighbor in NumPy and Its Equivalence to a Linear Layer

You are given a training set `X_train` (a NumPy array of shape `n × d`) with integer labels `y_train` (length `n`), and a test set `X_test` (shape `m × d`). Implement a 1-nearest-neighbor classifier: each test point receives the label of the training point closest to it in Euclidean (L2) distance. The constraint is that the computation must be **fully vectorized**: no Python loops over test points, training points or dimensions; only NumPy array operations. Then show that the same classifier can be written as a single linear layer `W x + b` followed by an argmax, and explain why. ### Clarifying Questions - How large are `n`, `m` and `d`? Does an `m × n` matrix of distances fit in memory? - How should ties (two training points at exactly the same distance) be broken? - Are the features on very different scales, or have very large magnitudes? - Is the output only the predicted labels, or also the distances or indices of the neighbors? ### Part 1 — Vectorized 1-nearest-neighbor Compute all test-to-train distances at once and return the predicted labels. ```hint Expand the square Write the squared distance between two vectors as a sum of three terms, and look for which of them can be computed for all pairs with one matrix multiplication. ``` #### What This Part Should Cover - A correct, loop-free distance computation with the right broadcasting shapes - Why squared distances suffice for choosing the nearest neighbor - Memory use of the intermediate arrays, and how to process test points in chunks - Numerical issues with the vectorized formula, and how to guard against them ### Part 2 — The same classifier as a linear layer Show that finding the nearest training point is equivalent to taking the argmax of `W x + b` for a suitable matrix `W` and vector `b` built from `X_train`. Give `W` and `b` explicitly, implement prediction that way, and explain what this equivalence is useful for. ```hint Drop what does not depend on the candidate In the expanded squared distance, one term is the same for every training point when the test point is fixed; see what remains after removing it. ``` #### What This Part Should Cover - A derivation of `W` and `b` from the expanded distance - Correct handling of the sign change from argmin to argmax - A check that both implementations agree - What the linear-layer view enables in practice ### What a Strong Answer Covers - Correct, fully vectorized NumPy code with shapes stated at each step - A clean derivation connecting nearest-neighbor search to a linear layer - Awareness of the memory footprint and of floating-point cancellation - Tests against a slow but obviously correct reference implementation ### Follow-up Questions - How would you return the `k` nearest neighbors and take a majority vote, still without loops? - The training set has 100 million points. How would you find nearest neighbors fast enough? - How does the linear-layer view change if you use cosine similarity instead of L2 distance? - How would you run this on a GPU framework instead of NumPy?