Scikit-Learn自定义插补器：用均值±标准差随机值填充NaN

我想要创建一个Custom Imputer，用NaN所在列的均值减标准差（mean - std）到均值加标准差（mean + std）范围内的随机值替换数据中的NaN值。以下是我目前编写的插补器代码：

class GroupImputer(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        X = check_array(X, force_all_finite=False)
        self.means = np.nanmean(X, axis=0)
        self.stds = np.nanstd(X, axis=0)
        return self
    def transform(self, X, y=None):
        check_is_fitted(self, 'means')
        check_is_fitted(self, 'stds')
        X = check_array(X, force_all_finite=False)
        # how do i apply to each row of the data?
        return 0

其中self.means存储了各列的均值列表，self.stds存储了各列的标准差列表。请问如何为数据中每一行的NaN值填充对应列mean - std到mean + std之间的随机值？是否需要遍历数据（如for row in X:）并根据列索引选取对应的均值和标准差？还是有更合适的实现方法？

最佳实现方案：向量化操作（强烈推荐） #

完全不需要逐行遍历！利用NumPy的向量化特性可以高效完成这个任务，比循环遍历快得多，尤其是处理大型数据集时。

思路很清晰：

1. 生成一个和输入X形状完全匹配的随机数组，每个元素都来自对应列的[mean-std, mean+std]区间
2. 用这个随机数组批量替换X中的NaN值

修改后的完整GroupImputer代码如下：

import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.utils.validation import check_array, check_is_fitted

class GroupImputer(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        X = check_array(X, force_all_finite=False)
        self.means = np.nanmean(X, axis=0)
        self.stds = np.nanstd(X, axis=0)
        return self
    
    def transform(self, X, y=None):
        check_is_fitted(self, ['means', 'stds'])
        # 加copy=True避免修改原始输入数据，符合Scikit-learn的API规范
        X = check_array(X, force_all_finite=False, copy=True)
        
        # 创建对应列的随机值数组：自动广播匹配X的形状
        random_vals = np.random.uniform(
            low=self.means - self.stds,
            high=self.means + self.stds,
            size=X.shape
        )
        
        # 批量替换NaN值
        np.where(np.isnan(X), random_vals, X, out=X)
        
        return X

关键细节说明： #

• copy=True：确保我们操作的是原数据的副本，不会污染输入的原始数据，这是Scikit-learn组件的标准行为
• np.random.uniform：支持传入数组作为low和high参数，NumPy会自动广播这些一维数组到X的二维形状，生成每列对应区间的随机值
• np.where：一次性完成所有NaN的替换操作，效率远高于逐元素循环

如果一定要用遍历实现（不推荐） #

虽然不建议，但如果你想通过遍历理解底层逻辑，也可以这样写（仅适合小数据集，效率极低）：

def transform(self, X, y=None):
    check_is_fitted(self, ['means', 'stds'])
    X = check_array(X, force_all_finite=False, copy=True)
    
    n_rows, n_cols = X.shape
    for i in range(n_rows):
        for j in range(n_cols):
            if np.isnan(X[i, j]):
                # 生成对应列的区间随机值
                low = self.means[j] - self.stds[j]
                high = self.means[j] + self.stds[j]
                X[i, j] = np.random.uniform(low, high)
    return X

这种方法在数据量大的时候会非常慢，所以优先选择向量化的实现方式。

内容的提问来源于stack exchange，提问作者Kevin