Implement 1NN with NumPy

Problem Statement

Implement a 1-Nearest Neighbor (OneNearestNeighbor) classifier in Python using NumPy. Your implementation must be fully vectorized—avoid explicit Python loops over test or training samples during prediction.

Implement the following class interface:

import numpy as np

class OneNearestNeighbor:
    def __init__(self, metric: str = "euclidean"):
        """
        metric: Supports 'euclidean' and 'cosine'.
        """
        pass

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        """
        Store training features and labels.
        X: shape (N, D), float64 array of training vectors.
        y: shape (N,), int64 array of training labels.
        """
        pass

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Predict labels for query samples.
        X: shape (M, D), float64 array of query vectors.
        Returns: shape (M,), int64 array of predicted labels.
        """
        pass

Metric Definitions

• Euclidean distance: For query vector u and training vector v, dist(u, v) = sqrt(sum((u - v)^2)). Note: minimizing squared Euclidean distance is equivalent to minimizing Euclidean distance.
• Cosine distance: dist(u, v) = 1 - (u . v) / (||u||_2 * ||v||_2). Assume vectors have non-zero norms.

Tie-breaking Rule

If a query point has equal minimum distance to multiple training points with differing labels, break ties by selecting the label associated with the smallest training index (i.e., the first point encountered in X_train).

Constraints

• 1 <= N, M <= 10_000
• 1 <= D <= 1_000
• N * M <= 50_000_000 (fits in memory if distance matrix is (M, N) float64, ~400 MB max).
• Do not use 3D broadcasting of shape (M, N, D) as it can trigger Out-Of-Memory errors for large D.
• Do not use scikit-learn or scipy.spatial.

Examples

Example 1: Euclidean Distance

X_train = np.array([[1.0, 1.0], [2.0, 2.0], [5.0, 5.0]])
y_train = np.array([0, 1, 2])
X_test = np.array([[1.1, 0.9], [4.5, 4.8]])

clf = OneNearestNeighbor(metric="euclidean")
clf.fit(X_train, y_train)
preds = clf.predict(X_test)
# Output: array([0, 2])
# Explanation: [1.1, 0.9] is closest to [1.0, 1.0] (label 0);
# [4.5, 4.8] is closest to [5.0, 5.0] (label 2).

Example 2: Tie-Breaking

X_train = np.array([[1.0, 0.0], [0.0, 1.0]])
y_train = np.array([10, 20])
X_test = np.array([[0.5, 0.5]])

clf = OneNearestNeighbor(metric="euclidean")
clf.fit(X_train, y_train)
preds = clf.predict(X_test)
# Output: array([10])
# Explanation: Distances to both points are equal; training index 0 is chosen.

Edge Cases

• M = 1 or N = 1.
• Large D where naive (M, 1, D) - (1, N, D) subtraction exceeds available RAM.