Implement a 1-Nearest-Neighbor Classifier with a Manhattan-Distance Follow-Up
Company: OpenAI
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: hard
Interview Round: Technical Screen
Implement a 1-nearest-neighbor (1-NN) classifier. You are given training points, each a vector of `d` numeric features with a class label, and a set of query points with the same `d` features. For every query, predict the label of the training point closest to it. Start with Euclidean distance; the follow-up asks you to support Manhattan distance as well.
### Clarifying Questions
- May I use NumPy, or should the solution use only the standard library?
- When two training points are exactly equally close to a query, which label should win?
- Roughly how many training points, query points and features are there? This decides whether computing every distance at once fits in memory.
- Should features be rescaled before computing distances, or used as given?
### Part 1 — Euclidean nearest neighbor
Write `predict(X_train, y_train, X_query)` that returns one predicted label per query point, using Euclidean distance.
```hint Share the work across pairs
Start from a correct double loop, then look for the parts of the distance computation that can be shared across all query-training pairs, and check what each version costs in memory.
```
#### What This Part Should Cover
- A correct nearest-neighbor search and label lookup with a stated tie rule.
- A vectorized distance computation and its memory footprint as the numbers of queries and training points grow.
- Time and space complexity of training and of prediction.
### Part 2 — Follow-up: Manhattan distance
Add support for Manhattan distance, the sum of absolute coordinate differences, selectable by a parameter. Explain what changes in your implementation and whether the switch can change the predictions.
```hint Check which speed-ups survive
Ask whether every trick you used to speed up the Euclidean version still applies once the distance is a sum of absolute values.
```
#### What This Part Should Cover
- Switching metrics without duplicating the search logic.
- Which speed-ups from Part 1 carry over to the new metric and which do not, with the memory consequences.
- A concrete case where the two metrics pick different neighbors, and the role of feature scaling.
### What a Strong Answer Covers
- A correct, tested implementation with small hand-checked examples, including a tie.
- Reasoning about when brute force is acceptable and what index structures or approximate search would replace it at scale.
- Awareness of 1-NN's sensitivity to noisy labels, irrelevant features and feature scale.
- Clean code: input validation, consistent shapes, and a single place where the metric is chosen.
### Follow-up Questions
- How would you extend this to k nearest neighbors with majority voting, and how would you break voting ties?
- With millions of training points and a latency budget per query, what would you change?
- How does high dimensionality affect nearest-neighbor methods, and does the choice between the two metrics matter more there?
- How would you choose between the two metrics empirically on a given dataset?
Overview: Implement a 1-nearest-neighbor classifier that labels each query point with the label of its closest training point under Euclidean distance, then extend it to Manhattan distance. Tests vectorized distance computation under memory limits, tie handling and how the choice of metric changes predictions.
Community answers
Answer by panvaf
import numpy as np
def softmax(z, axis=-1):
z = z - np.max(z, axis=axis, keepdims=True) # for numerical stability
exp_z = np.exp(z)
return exp_z / np.sum(exp_z, axis=axis, keepdims=True)
def relu(z):
return np.maximum(z, 0.0)
---------- Part 1: vectorized 1-NN (squared L2) ----------
def one_nn_predict(X_train, y_train, X_query):
# ||q - x||^2 = ||q||^2 - 2 q.x + ||x||^2
# ||q||^2 is the same for every training point, so it can't change the argmin -> skip it.
train_sq_norms = np.sum(X_train ** 2, axis=1) # (n,)
cross = X_query @ X_train.T # (m, n)
scores = train_sq_norms[None, :] - 2 * cross # (m, n)
nearest = np.argmin(scores, axis=1) # (m,), argmin takes the first min on ties
return y_train[nearest]
---------- Part 2: 1-NN as a forward pass (Linear + softmax) ----------
def one_nn_network_forward(X_train, X_query, beta=1.0):
# One output unit per training example:
# logit_i = 2 x_i.q - ||x_i||^2 = ||q||^2 - ||q - x_i||^2
# so the largest logit belongs to the nearest neighbor.
W = 2 * X_train.T # (d, n)
b = -np.sum(X_train ** 2, axis=1) # (n,)
logits = X_query @ W + b # (m, n)
return softmax(beta * logits, axis=1) # (m, n); bigger beta -> closer to one-hot
---------- L1 version ----------
def one_nn_predict_l1(X_train, y_train, X_query):
diffs = X_query[:, None, :] - X_train[None, :, :] # (m, n, d)
l1_dists = np.sum(np.abs(diffs), axis=2) # (m, n)
nearest = np.argmin(l1_dists, axis=1)
return y_train[nearest]
def one_nn_l1_network_forward(X_train, X_query, beta=1.0):
# No abs