k-Nearest-Neighbors Classifier

k-Nearest-Neighbors is the simplest non-parametric classifier: to label a query point, look at the k closest labeled examples and take a vote. There is no training step beyond storing the data, which makes the whole problem about computing distances efficiently and resolving votes deterministically.

Task

Implement knn_predict(X_train, y_train, X_test, k). For each row of X_test:

1. Compute the Euclidean distance to every row of X_train.
2. Select the k training points with the smallest distances.
3. Return the majority label among those k neighbors.

Compute the pairwise distances with vectorized NumPy (broadcasting), not nested Python loops. Break ties deterministically: when several training points are equidistant, prefer the one with the smaller index; when two labels receive the same vote count, prefer the smaller label id. Return an integer array of shape (n_test,).

Example

import numpy as np

X_train = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [5.0, 5.0], [6.0, 5.0]])
y_train = np.array([0, 0, 0, 1, 1])
X_test = np.array([[0.5, 0.5], [5.5, 5.0]])

print(knn_predict(X_train, y_train, X_test, k=3))
# [0 1]

Constraints

• X_train has shape (n_train, d), X_test has shape (n_test, d), y_train has shape (n_train,).
• 1 <= k <= n_train; labels are non-negative integers.
• Use only NumPy; do not call scikit-learn, SciPy, or any distance/kNN helper.
• Distances must be computed in a vectorized manner (no explicit per-pair Python loop).

Approach

Build the full (n_test, n_train) distance matrix with broadcasting, so no Python loop touches individual pairs. For each query row, take the indices of the k smallest distances using a stable argsort (stable order makes equidistant points break ties toward the smaller training index). Gather those neighbors’ labels, tally them with np.bincount, and pick the label with the most votes; argmax returns the smallest index on a tie, which is exactly the smaller-label rule.

Solution

import numpy as np

def knn_predict(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    k: int,
) -> np.ndarray:
    """Classify each test point by majority vote of its k nearest neighbors."""
    X_train = np.asarray(X_train, dtype=float)
    X_test = np.asarray(X_test, dtype=float)
    y_train = np.asarray(y_train)

    # Squared Euclidean distances, shape (n_test, n_train), via the identity
    # ||a - b||^2 = ||a||^2 - 2 a.b + ||b||^2. Monotonic in true distance, so
    # skipping the sqrt does not change which neighbors are nearest.
    train_sq = np.sum(X_train ** 2, axis=1)             # (n_train,)
    test_sq = np.sum(X_test ** 2, axis=1)               # (n_test,)
    cross = X_test @ X_train.T                           # (n_test, n_train)
    dist_sq = test_sq[:, None] - 2.0 * cross + train_sq[None, :]

    # Stable sort -> ties in distance resolve to the smaller training index.
    order = np.argsort(dist_sq, axis=1, kind="stable")   # (n_test, n_train)
    knn_idx = order[:, :k]                                # (n_test, k)
    knn_labels = y_train[knn_idx]                         # (n_test, k)

    n_classes = int(y_train.max()) + 1
    preds = np.empty(X_test.shape[0], dtype=y_train.dtype)
    for i in range(knn_labels.shape[0]):
        votes = np.bincount(knn_labels[i], minlength=n_classes)
        # argmax returns the smallest index on a tie -> smaller label wins.
        preds[i] = np.argmax(votes)
    return preds

Walkthrough

Take the example with k=3.

• X_test[0] = [0.5, 0.5]. Squared distances to the five training points are [0.5, 0.5, 0.5, 40.5, 50.5]. The stable argsort keeps indices 0, 1, 2 first, so the three neighbors are the three cluster-0 points, labels [0, 0, 0]. bincount gives [3, 0], argmax is 0.
• X_test[1] = [5.5, 5.0]. Squared distances are [55.25, 45.25, 46.25, 0.25, 0.25]. The two closest are indices 3 and 4 (both 0.25, index 3 first by stability), then index 1. Labels are [1, 1, 0], bincount gives [1, 2], argmax is 1.

Result: [0, 1], matching the expected output.

Complexity & notes

• Time: O(n_test * n_train * d) to form the distance matrix, plus O(n_test * n_train * log n_train) for the per-row sort. Space: O(n_test * n_train) for the distance matrix.
• Using squared distance avoids an unnecessary sqrt over the whole matrix and never affects the ranking. Floating-point noise in the ||a||^2 - 2a.b + ||b||^2 form can produce tiny negatives for coincident points; that is harmless here because we only sort, but clip to 0 before any real sqrt.
• For a large n_train, np.argpartition(dist_sq, k, axis=1)[:, :k] finds the k smallest in O(n_train) instead of a full sort; you then need an explicit tie-break because partition is not stable.
• Both tie-breaks lean on NumPy’s determinism: kind="stable" for equidistant points and argmax’s smallest-index rule for equal vote counts. State these rules in the interview; a naive Counter.most_common is order-dependent and not reproducible.