k-Nearest Neighbors › Implementation in NumPy and scikit-learn

Michael BrenndoerferOctober 26, 202612 min read
Machine Learning from ScratchAll chapters
Part 7 of 15Supervised Learning Beyond Linear Models
Chapter 53k-Nearest Neighbors › Implementation in NumPy and scikit-learn
Reading section 7 of 10 · Implementation in NumPy and scikit-learn · 12 minSection 7 of 10 · 12 min

Section 7 of 10 · 12 min read

k-Nearest Neighbors › Implementation in NumPy and scikit-learnSection view · Sections in this chapter7 / 10 ↓
Reading tools Definitions, links & notes

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Notes & comments

We now turn the formulas into code using only NumPy primitives, then compare the result with scikit-learn on the same synthetic data. Writing it by hand makes every step auditable: distance computation, sorting, and voting or averaging. It is not intended to replace scikit-learn. Compare results with the same distance, weighting, and tie conventions, allowing for numerical tolerance. A remaining discrepancy calls for inspecting both implementations; a different tie rule can produce a different valid result.

We will deliberately use a one-dimensional, noisy regression problem first, because it is the easiest setting to visualize and to reason about. Then we build a classification example. The examples below also use explicit NumPy generator seeds.

PythonIn [12]
CollapseExpand
from sklearn.metrics import mean_squared_error

rng = np.random.default_rng(42)
X_reg = rng.uniform(-3.0, 3.0, size=(60, 1))
y_true_reg = np.sin(X_reg[:, 0])
y_reg = y_true_reg + rng.normal(0.0, 0.25, size=60)

X_grid = np.linspace(-3.0, 3.0, 300).reshape(-1, 1)
y_grid_true = np.sin(X_grid[:, 0])

# Compute all pairwise squared distances between query grid and training points
sq_dists = ((X_grid[:, None, :] - X_reg[None, :, :]) ** 2).sum(axis=2)
mean_sq_dist = float(sq_dists.mean())
ConsoleOut [13]
CollapseExpand
Mean squared query-to-training distance: 5.882

This average summarizes squared distances over every query-training pair in the chosen feature units. It reflects the spread of both sets, not just local support: even queries equal to training points can have a large all-pair average. Inspect nearest-neighbor distances and the training domain when assessing whether a particular query lacks nearby observations.

The majority-vote classifier used in the earlier figures is already defined above. The next cell defines the regression predictor and its weight helper, using the same distance calculation and stable neighbor ordering.

PythonIn [14]
CollapseExpand
def inverse_squared_weights(distances):
    """Normalized-equivalent 1/d² weights; exact matches receive all weight."""
    distances = np.asarray(distances, dtype=float)
    zero = distances == 0.0
    has_zero = zero.any(axis=1, keepdims=True)
    safe_distances = np.where(zero, 1.0, distances)
    scale = safe_distances.min(axis=1, keepdims=True)
    w = (scale / safe_distances) ** 2
    return np.where(has_zero, zero.astype(float), w)


def knn_regression_predict(X_train, y_train, X_query, k, weights="uniform"):
    """Brute-force kNN for regression.

    Parameters
    ----------
    X_train : array, shape (n, p)
    y_train : array, shape (n,)
    X_query : array, shape (m, p)
    k : int
        Number of neighbors to use.
    weights : {"uniform", "distance"}
        Weighting scheme for averaging neighbor targets.

    Returns
    -------
    y_pred : array, shape (m,)
    """
    X_train = np.asarray(X_train)
    X_query = np.asarray(X_query)
    y_train = np.asarray(y_train)
    if not isinstance(k, (int, np.integer)) or not 1 <= k <= len(X_train):
        raise ValueError(
            "k must be an integer between 1 and the training-set size"
        )

    # Pairwise distances (Euclidean)
    diffs = X_query[:, None, :] - X_train[None, :, :]
    dists = np.sqrt((diffs**2).sum(axis=2))

    # Indices of the k nearest neighbors, sorted by distance
    nn_idx = np.argsort(dists, axis=1, kind="stable")[:, :k]

    # Gather labels of neighbors and reduce
    neighbor_labels = y_train[nn_idx]

    if weights == "uniform":
        y_pred = neighbor_labels.mean(axis=1)
    elif weights == "distance":
        neighbor_dists = np.take_along_axis(dists, nn_idx, axis=1)
        w = inverse_squared_weights(neighbor_dists)
        y_pred = (w * neighbor_labels).sum(axis=1) / w.sum(axis=1)
    else:
        raise ValueError(f"Unknown weights: {weights}")

    return y_pred

We now use the noisy sine curve generated above to compare uniform and distance-weighted kNN across several kk values. The following block calculates predictions for both weighting schemes and the uniform-weighting errors against the known sine signal.

PythonIn [15]
CollapseExpand
ks = [1, 5, 15]
uniform_preds = {
    k: knn_regression_predict(X_reg, y_reg, X_grid, k, "uniform") for k in ks
}
weighted_preds = {
    k: knn_regression_predict(X_reg, y_reg, X_grid, k, "distance") for k in ks
}
rmse_true = {
    k: float(np.sqrt(mean_squared_error(y_grid_true, predictions)))
    for k, predictions in uniform_preds.items()
}
ConsoleOut [16]
CollapseExpand
k =  1   RMSE vs true signal = 0.216
k =  5   RMSE vs true signal = 0.146
k = 15   RMSE vs true signal = 0.207

These errors compare predictions with the known noiseless sine signal on the display grid. Lower values mean a closer approximation to that particular signal; these are teaching metrics, not a held-out estimate for a new population.

::: {layout-ncol=3}

VisualizationOut [17]
CollapseExpand
One regression plot with noisy training points, the known sine signal and uniform kNN predictions for k=1.
Uniform kNN regression with k=1 on the noisy sine observations. Each query snaps to the label of the single nearest training point, so the prediction is a stair-step curve that follows training noise rather than the smooth sine signal.
VisualizationOut [18]
CollapseExpand
One regression plot with noisy training points, the known sine signal and uniform kNN predictions for k=5.
Uniform kNN regression with k=5 on the noisy sine observations. Averaging five targets smooths the stair-step noise of k=1 while still tracking the broad shape of the sine curve.
VisualizationOut [19]
CollapseExpand
One regression plot with noisy training points, the known sine signal and uniform kNN predictions for k=15.
Uniform kNN regression with k=15 on the noisy sine observations. Averaging fifteen targets flattens local variation over a larger neighborhood, trading resolution for stability.

:::

The first regression figure shows the stair-step behavior of k=1k=1: each query snaps to the label of the single nearest training point. Increasing kk averages more targets and changes the balance between noise and smoothing bias. Here k=15k=15 flattens local variation over a larger neighborhood; the computed RMSEs show how that tradeoff affects approximation of the known sine signal. This is the same tension we saw earlier in the 2D classification picture, now in regression form: the winning value of kk balances the two failure modes against each other.

The following block fits the scikit-learn regressors and records the maximum absolute differences against our NumPy predictions, confirming that our brute-force predictions match scikit-learn's implementation exactly and verifying distance-weighted prediction. We pass inverse_squared_weights as a callable so both implementations use the same weighting rule; scikit-learn’s built-in weights="distance" instead uses inverse distance.

PythonIn [20]
CollapseExpand
# Compare uniform kNN regression with scikit-learn
knn_sk = KNeighborsRegressor(n_neighbors=5, weights="uniform")
knn_sk.fit(X_reg, y_reg)
y_sk = knn_sk.predict(X_grid)
y_ours = uniform_preds[5]
max_diff_uniform = float(np.max(np.abs(y_sk - y_ours)))

# Distance-weighted comparison
knn_sk_w = KNeighborsRegressor(n_neighbors=5, weights=inverse_squared_weights)
knn_sk_w.fit(X_reg, y_reg)
y_sk_w = knn_sk_w.predict(X_grid)
y_ours_w = weighted_preds[5]
max_diff_weighted = float(np.max(np.abs(y_sk_w - y_ours_w)))
rmse_uniform_5 = float(np.sqrt(mean_squared_error(y_grid_true, y_ours)))
rmse_weighted_5 = float(np.sqrt(mean_squared_error(y_grid_true, y_ours_w)))
ConsoleOut [21]
CollapseExpand
Max |uniform - sklearn|  = 0.000e+00
Max |weighted - sklearn| = 0.000e+00
Uniform k=5 RMSE vs true signal: 0.146
Weighted k=5 RMSE vs true signal: 0.169

Both differences are at floating-point noise level, confirming that the primitive operations in our NumPy implementation correspond to scikit-learn's predictions. The hand-written version and the library are doing the same arithmetic; writing it out simply makes the steps visible.

Distance weighting changes the regression curve without changing the neighbor set. In this seeded comparison, inverse-squared weighting has RMSE 0.169 against the known sine signal, compared with 0.146 for uniform weighting; giving closer observations more influence does not improve this whole-grid score. The next figure overlays the uniform and inverse-squared weighted predictions at k=5k=5 on the same noisy sine data, using the RMSE values computed above.

VisualizationOut [22]
CollapseExpand
Uniform and inverse-squared weighted kNN regression curves at k=5.
Uniform and inverse-squared distance-weighted kNN regression at k=5 on the noisy sine data. Both curves use the same neighbors, so any gap shows how much the weighting changes the local average near unevenly spaced observations. The RMSE values in the legend compare each curve against the known sine signal.

Now classification, with training-only scaling. We split the raw data into training and validation portions, fit the scaler on the training portion only, and compare several values of kk on the validation portion. Because this comparison can guide the choice of kk, these held-out observations serve as validation data; an additional untouched test set would be needed to assess the selected model.

PythonIn [23]
CollapseExpand
# Two moderately overlapping 2-class problem
X_cls = np.vstack(
    [
        rng.normal(loc=(1.0, 1.0), scale=1.0, size=(100, 2)),
        rng.normal(loc=(2.5, 2.5), scale=1.0, size=(100, 2)),
    ]
)
y_cls = np.array([0] * 100 + [1] * 100)

# Split first, then fit the scaler on the training portion only
X_tr_raw, X_te_raw, y_tr, y_te = train_test_split(
    X_cls,
    y_cls,
    test_size=0.25,
    random_state=0,
    stratify=y_cls,
)
scaler = StandardScaler().fit(X_tr_raw)
X_tr = scaler.transform(X_tr_raw)
X_te = scaler.transform(X_te_raw)

# Compare several k values with scikit-learn
k_list = [1, 3, 5, 9, 15, 25]
acc_scaled, acc_unscaled = [], []
for k in k_list:
    m_scaled = KNeighborsClassifier(n_neighbors=k).fit(X_tr, y_tr)
    m_raw = KNeighborsClassifier(n_neighbors=k).fit(X_tr_raw, y_tr)
    acc_scaled.append(accuracy_score(y_te, m_scaled.predict(X_te)))
    acc_unscaled.append(accuracy_score(y_te, m_raw.predict(X_te_raw)))

# Baseline learned from training labels, evaluated on the same validation rows.
majority_label = np.argmax(np.bincount(y_tr))
majority_accuracy = float(np.mean(y_te == majority_label))
ConsoleOut [24]
CollapseExpand
k =  1   scaled accuracy = 0.740   raw (unscaled) accuracy = 0.740
k =  3   scaled accuracy = 0.800   raw (unscaled) accuracy = 0.800
k =  5   scaled accuracy = 0.780   raw (unscaled) accuracy = 0.780
k =  9   scaled accuracy = 0.820   raw (unscaled) accuracy = 0.820
k = 15   scaled accuracy = 0.860   raw (unscaled) accuracy = 0.860
k = 25   scaled accuracy = 0.900   raw (unscaled) accuracy = 0.880

On this particular synthetic data, standardization changes little because the two features have similar variances, so scaling does not substantially change which points are nearest. The printed accuracies show which kk performs best on this split and how scaled and raw features compare; read them against the majority-class baseline of 0.50. These results do not demonstrate the large-kk underfitting limit; that limit concerns neighborhoods approaching the full training-set size. When feature units differ, scaling can change neighborhoods substantially. Whether this improves prediction depends on the target and should be checked through validation; the earlier visualization isolates the geometric change.

Validation accuracy across the candidate neighborhood sizes makes this comparison visible. The next figure plots the already computed accuracies for scaled and raw features against the training-majority baseline. The best value on this validation split is a tuning result, not an untouched test estimate for the selected model.

VisualizationOut [25]
CollapseExpand
Validation accuracy versus k for scaled and raw features, with a dashed training-majority baseline.
Validation accuracy across candidate neighborhood sizes on the synthetic two-class split. Both curves use the same training and validation rows, and the dashed line shows the computed training-majority baseline. Selecting k from these curves requires another untouched test set for final evaluation.

Key Methods

This section introduces the following fitting and prediction methods, plus methods for inspecting neighbors:

  • fit(X, y): Stores the training examples and labels and may build a search index; neighbor lookup and prediction happen for each query. The kneighbors method inspects the selected neighbors separately from prediction.
  • predict(X): Returns a class label or numeric prediction for each query in X, using the configured voting or averaging weights.
  • predict_proba(X): Returns class-support estimates normalized to sum to one (classification only). Uniform weights give neighbor-count fractions; distance weights give normalized weighted totals. Assess their calibration on held-out observations.
  • kneighbors(X): Returns the distances and indices of the kk nearest training points for each query, useful for inspecting neighborhoods and identifying queries with weak local support.

Key Parameters

Below are some of the main parameters that affect how kNN works and performs:

  • n_neighbors / kk: Number of neighbors used for each prediction. Small values can follow individual observations closely and can be sensitive to changes in them. Larger values average a broader neighborhood and can smooth away local signal. The effect on bias and variance depends on the data and target, and on the weights. Choose from a validation grid such as {1,3,5,9,15,25}\{1, 3, 5, 9, 15, 25\}, evaluated under the same protocol for all compared models.
  • weights: Uniform weights treat all neighbors equally. Distance weighting lets closer neighbors count more; this chapter uses inverse-squared distance, while scikit-learn's weights="distance" uses inverse distance.
  • metric / p: Distance used to define the neighborhood. Euclidean (metric="minkowski", p=2) is the default; Manhattan (p=1) sums absolute coordinate differences and can select different neighbors.
  • algorithm: Search strategy. "auto" picks a strategy in scikit-learn; the NumPy code scans and sorts all training distances.

Comments

Add to the conversation

A question, a correction, or an example from your own work.

No comments yet. Be the first to share your thoughts!

Newsletter

Stay up to date

Get new articles, book updates, and notes on AI and technology.

Unsubscribe at any time. Privacy policy