I have a dataset of 70 columns that have missing values. Each column has a few columns (3-5) that it is significantly more correlated than the others but each column's correlated columns are very different from other columns. I would like to perform a tweaked version of KNN imputation and before I start writing this from scratch, I'd like to know if there's something similar out there, So, I don't go about reinventing the wheel.

On initialization, a correlation matrix is obtained between all the variables. For every missing cell that needs to be imputed, the algorithm retrieves from the correlation matrix the top n highest correlated columns and uses only those N columns to select the k nearest neighbors and use the mean value of those neighbors.

Thanks and apologies if I am not very clear with my question. I'll try to clarify if you have any doubts.

  • $\begingroup$ Unfortunately I am not aware of such a package already existing but I think it is an interesting idea. If computationally it is not to heavy maybe you can also test to calculate the distance by weighting them with their correlation. $\endgroup$
    – Viktor
    Commented Jun 21, 2019 at 17:24
  • $\begingroup$ Use a for-loop on each column - 1. to select your N Columns and 2. use sklearn.impute.KNNImputer as suggested by @Brian to predict and fill the missing values $\endgroup$
    – 10xAI
    Commented May 22, 2020 at 11:27

1 Answer 1


You can extend scikit-learn's sklearn.impute.KNNImputer. The metric parameter can accept the custom function you described.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.