Problems
Loading...
1 / 1

KNN Distance + Neighbor Lookup

Classic ML
Medium

Given training data, test data, and k, compute pairwise Euclidean distances and return the k nearest neighbor indices for each test point.

This is the core computational step in k-Nearest Neighbors (KNN) algorithm. Understanding efficient distance computation and neighbor lookup is crucial for implementing KNN classifiers and regressors.

Euclidean Distance Formula:

d(x,y)=∑i=1d(xi−yi)2d(\mathbf{x}, \mathbf{y}) = \sqrt{ \sum_{i=1}^{d} (x_i - y_i)^{2} }

Where x,y∈Rd are feature vectors and d is the number of dimensions.

Function Arguments

  • X_train: array-like, shape (n_train, d) or (n_train,) - Training features
  • X_test: array-like, shape (n_test, d) or (n_test,) - Test features
  • k: int - Number of nearest neighbors to find
Loading visualization...

Examples

Input: X_train=[1,3,5], X_test=[2], k=2

Output: [[0,1]]

Input: X_train=[[0,0],[1,1],[2,2]], X_test=[[0.5,0.5]], k=2

Output: [[0,1]]

Hint 1

Use broadcasting to compute all pairwise distances: X_test[:, np.newaxis, :] - X_train[np.newaxis, :, :].

Hint 2

Use np.argsort() to get indices of sorted distances, then take the first k.

Hint 3

Handle 1D arrays by reshaping to 2D: X.reshape(-1, 1).

Requirements

  • Return numpy array of shape (n_test, k) with integer dtype
  • Each row contains indices of k nearest training points
  • Use vectorized distance computation (broadcasting)
  • Handle k larger than training set size (pad with -1)
  • Support both 1D and multi-dimensional features
  • Sort neighbors by distance (closest first)

Constraints

  • n_train, n_test ≤ 1000; d ≤ 100; NumPy only
  • Time limit: 200ms; Memory: 128MB
Try Similar Problems
Euclidean DistanceManhattan DistanceCosine SimilarityK Means AssignmentK Means Centroid Update
PrevNext