ClusterImputer

A scikit-learn compatible Cluster-Aware Tabular Imputer for mixed numerical and categorical datasets, published in Nature Scientific Reports.

• Paper: Cluster-aware imputation: a hybrid XGBoost-based model for missing data
• Journal: Scientific Reports (Springer Nature), Volume 16, Article 70170 (2026)
• DOI: 10.1038/s41598-026-70170-9

Overview

ClusterImputer (also available as ClusteringImputer) addresses the limitation of traditional MICE global regression models by learning local correlation structures within tabular data:

1. Missingness-Aware Chaining: Features are prioritized and imputed in ascending order of missingness (least missing to most missing).
2. Subspace Clustering: For each feature, complete rows are projected into PCA space and partitioned using K-Means with automatic silhouette score optimization ($k$).
3. Parallel Localized Ensembles: High-speed XGBoost regressors and classifiers are fitted per cluster in parallel.
4. Iterative MICE Convergence: Chained updates repeat iteratively until numerical tolerance or maximum iterations are satisfied.
5. Linear Computational Scaling: Benchmarked across 97 UCI datasets, demonstrating asymptotic linearity $\mathcal{O}(N)$ in computational complexity compared to KNN's quadratic $\mathcal{O}(N^2)$ scaling.

Installation

pip install cluster-imputer

Dependencies

• numpy >= 1.20.0
• pandas >= 1.3.0
• scikit-learn >= 1.0.0
• xgboost >= 1.5.0
• joblib >= 1.1.0

Quick Start

import numpy as np
import pandas as pd
from cluster_imputer import ClusterImputer

# 1. Prepare sample data with missing entries
df = pd.DataFrame({
    "feature_1": [1.2, np.nan, 3.4, 4.1, 5.0, 2.8, np.nan, 3.9],
    "feature_2": [10.5, 20.1, np.nan, 40.2, 51.0, 22.3, 31.1, 41.5],
    "category":  ["A", "B", "A", np.nan, "B", "A", "B", np.nan],
})

# 2. Initialize the imputer
imputer = ClusterImputer(
    n_clusters='auto',   # Automatically finds optimal k via silhouette score
    max_k=5,
    max_iter=5,
    n_jobs=-1
)

# 3. Fit and transform
df_imputed = imputer.fit_transform(df)
print(df_imputed)

Scikit-Learn Pipeline Integration

ClusterImputer inherits from BaseEstimator and TransformerMixin, making it seamlessly compatible with Pipeline, ColumnTransformer, and cross-validation:

from sklearn.pipeline import Pipeline
from sklearn.ensemble import RandomForestClassifier
from cluster_imputer import ClusterImputer

pipe = Pipeline([
    ("imputer", ClusterImputer(n_clusters=3, max_iter=5, n_jobs=-1)),
    ("classifier", RandomForestClassifier(n_estimators=100, random_state=42))
])

pipe.fit(X_train, y_train)
predictions = pipe.predict(X_test)

Hyperparameter Selection

To estimate recommended parameters based on dataset dimensionality and missingness ratio, use suggest_imputer_params:

from cluster_imputer import suggest_imputer_params

params = suggest_imputer_params(df)
print("Recommended parameters:", params)

imputer = ClusterImputer(**params)
df_clean = imputer.fit_transform(df)

Parameters and API Reference

ClusterImputer / ClusteringImputer

Methods

• fit(X, y=None): Learns clustering spaces and per-cluster XGBoost models on observed data.
• transform(X): Imputes missing values in X and returns a completed DataFrame or numpy ndarray.
• fit_transform(X, y=None): Fits and transforms the data in a single call.

Citation

If you use ClusterImputer in your research, please cite our article in Scientific Reports:

@article{ismail2026clusteraware,
  title = {Cluster-aware imputation: a hybrid XGBoost-based model for missing data},
  author = {Ismail, Mahmoud M. and Emad, Mohamed and Mohamed, Mai},
  journal = {Scientific Reports},
  volume = {16},
  number = {1},
  pages = {70170},
  year = {2026},
  publisher = {Springer Nature},
  doi = {10.1038/s41598-026-70170-9},
  url = {https://doi.org/10.1038/s41598-026-70170-9}
}

License

This project is licensed under the MIT License - see the LICENSE file for details.