KNN Imputation

September 2026 · 11 min read

KNN imputation fills a missing feature by borrowing observed values from similar training rows. It replaces a single global mean with a local average. That can help when the observed features identify useful neighbors; it can also perform worse than a simple fill when distance is misleading.

Follow one cell first. Our query row is [1, 1, missing]. We want its third feature. The first two features are numeric measurements already on comparable scales in this constructed example.

Choose donors for this particular gap

A donor must have an observed value in the column we are filling. It also needs at least one observed feature in common with the query so that a distance can be calculated. A donor does not need to be complete in every column.

Here A and B are the two closest usable donors, carrying third-feature values 10 and 12. With k = 2 and equal weights, the fill is their mean: 11. The missing cell never participates in the distance calculation.

This donor search is specific to the missing column. If another feature in the same query is missing, its eligible donors can differ. Rows with an observed third feature might be missing the fourth, and vice versa.

Measure distance using the features both rows actually have

Ordinary Euclidean distance adds squared differences over every feature and takes a square root. Here some differences are unavailable. Scikit-learn’s nan_euclidean rule keeps the shared coordinates and increases their weight to compensate for the reduced count.

Let $D$ be the total number of features, $S$ the set of coordinates observed in both rows, and $m$ the number of coordinates in $S$. For query $q$ and donor $x$, when $m>0$:

For donor B, $D=3$ and $m=2$. Only the second coordinate differs, by 1, so the distance is $\sqrt{(3/2)(0^2+1^2)}\approx1.225$. With no shared coordinates, the distance is undefined, not zero.

The factor $D/m$ adjusts the scale; it cannot invent evidence about unobserved coordinates. Two rows matching on one feature can have distance zero while disagreeing on everything hidden. This is one reason to examine overlap as well as distance.

Average the donors, and specify the fallback

Uniform weighting gives the selected donors equal influence. Distance weighting gives closer donors more influence, using inverse distance. In scikit-learn, if selected donors have distance zero, only those zero-distance donors receive weight. In our example, distance weighting therefore fills 10, even with k = 2.

Fit the distance rule and donor pool on training data

A feature measured in thousands can dominate one measured in fractions. Scaling before neighbor search changes which rows count as similar. StandardScaler estimates each column’s statistics from its observed training entries and preserves NaNs during transformation, so the imputer can still recognize gaps.

Scaling is a modeling choice, not a guarantee of useful distance. Irrelevant dimensions, duplicated measurements, arbitrary category codes, and many sparse one-hot columns can all distort the neighborhood. Use features and a distance rule suited to the data.

import numpy as np
from sklearn.impute import KNNImputer

train = np.array([
    [1, 1, 10], [1, 2, 12], [9, 9, 99],
    [np.nan, np.nan, 50], [1, 1, np.nan],
])
query = np.array([[1, 1, np.nan]])
imputer = KNNImputer(n_neighbors=2, weights="uniform").fit(train)
filled = imputer.transform(query)
assert np.allclose(filled, [[1, 1, 11]])
print(filled)

# No shared coordinates: use each training column's observed mean.
no_information = imputer.transform([[np.nan, np.nan, np.nan]])
assert np.allclose(no_information, [[3, 3.25, 42.75]])

from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

preprocess = make_pipeline(
    StandardScaler(),
    KNNImputer(n_neighbors=2, weights="uniform"),
).fit(train)
scaled_filled = preprocess.transform(query)
# The output is in standardized units. Invert the scaler only if
# you need to inspect the filled feature values in original units.
original_units = preprocess.named_steps["standardscaler"].inverse_transform(
    scaled_filled
)
assert np.allclose(original_units, [[1, 1, 11]])

Decide whether the extra work earns its keep

Compare the complete predictive pipeline with a median or mean baseline on the same validation splits. Tune k and weighting using those splits. A better reconstruction of masked feature values does not necessarily improve the final prediction task.

If you artificially hide known values, state how you chose the mask. Uniform random masking tests a different situation from a sensor that fails mostly at extreme readings. The missingness mechanism affects what that experiment can tell you; KNN itself does not establish MAR or solve MNAR.

For $M$ query rows, $N$ training donors, and $D$ features, a brute-force distance calculation costs roughly $O(MND)$. Imputing the training set can require roughly $O(N^2D)$ distance work. Implementations can reuse distances across missing columns and process them in chunks; neighbor selection adds work, and the donor pool still needs storage. Measure memory and serving latency as well as predictive quality.

C++: one cell, a frozen donor pool, and a clear fallback

This compact implementation uses uniform weights and the distance rule above. It does not mutate the training rows or feed newly filled values back into the donor pool. If no distance is available, it uses the observed column mean; if the target column has no observations, it requires an explicit decision by throwing an exception.

#include <algorithm>
#include <cmath>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <utility>
#include <vector>

using Row = std::vector<double>;

double fillCell(const std::vector<Row>& train, const Row& query,
                std::size_t column, std::size_t k) {
    if (train.empty() || query.empty() || column >= query.size() || k == 0)
        throw std::invalid_argument("nonempty input, valid column and k required");
    auto validate = [&](const Row& row) {
        if (row.size() != query.size())
            throw std::invalid_argument("feature counts differ");
        for (double x : row)
            if (std::isinf(x)) throw std::invalid_argument("infinite feature");
    };
    validate(query);
    for (const auto& row : train) validate(row);
    if (!std::isnan(query[column])) return query[column];

    std::vector<std::pair<double, std::size_t>> candidates;
    double columnSum = 0;
    std::size_t observed = 0;
    for (std::size_t i = 0; i < train.size(); ++i) {
        const auto& row = train[i];
        if (std::isnan(row[column])) continue;
        columnSum += row[column];
        ++observed;
        double squared = 0;
        std::size_t shared = 0;
        for (std::size_t j = 0; j < query.size(); ++j) {
            if (std::isnan(query[j]) || std::isnan(row[j])) continue;
            const double difference = query[j] - row[j];
            squared += difference * difference;
            ++shared;
        }
        if (shared)
            candidates.push_back({std::sqrt(
                static_cast<double>(query.size()) / shared * squared), i});
    }
    if (!observed) throw std::domain_error("target column has no observations");
    if (candidates.empty()) return columnSum / observed;

    // Equal distances are resolved by donor row index, not donor target value.
    std::sort(candidates.begin(), candidates.end());
    const std::size_t count = std::min(k, candidates.size());
    double sum = 0;
    for (std::size_t i = 0; i < count; ++i)
        sum += train[candidates[i].second][column];
    return sum / count;
}

int main() {
    const double NA = std::numeric_limits<double>::quiet_NaN();
    const std::vector<Row> train = {
        {1, 1, 10}, {1, 2, 12}, {9, 9, 99}, {NA, NA, 50}, {1, 1, NA}
    };
    std::cout << "nearest-two fill: " << fillCell(train, {1, 1, NA}, 2, 2) << "\n";
    std::cout << "no-overlap fallback: " << fillCell(train, {NA, NA, NA}, 2, 2) << "\n";
}

nearest-two fill: 11
no-overlap fallback: 42.75

The example is for moderate finite numerical inputs. For a production pipeline, use the library implementation and check its options, numerical behavior, and feature schema.

References and further reading

• Scikit-learn: KNNImputer — neighbor counts, weighting choices, stored training data, and empty-feature behavior.
• Scikit-learn: missing-aware Euclidean distance — the total-feature/shared-feature scaling and undefined-distance case.
• Scikit-learn: nearest-neighbor imputation guide — per-feature donors and fallback behavior.
• Scikit-learn: StandardScaler — training statistics and preservation of missing values.