如何基于分组字段随机采样填充Pandas DataFrame列缺失值

高效实现方案 #

核心思路是规避Python层面的循环，利用pandas原生优化的关联、分组操作实现向量化计算，性能比逐组遍历的方案高1~2个数量级，完全适配大规模数据集。

方案1：左连接+分组采样（通用性最强）

import pandas as pd
import numpy as np

# 示例数据生成
source = pd.DataFrame({'age':5*[21],
                       'location':[0,0,1,1,1],
                       'x':[1,2,3,4,4]})
target = pd.DataFrame({'age':5*[21],
                       'location':[0,0,0,1,2],
                       'x':5*[np.nan]})

# 核心填充逻辑
# 1. 提取待填充的缺失行，添加原始行索引用于后续回写
target_null = target[target['x'].isna()].reset_index(names='orig_idx')

# 2. 左连接匹配source中同age、location组合的所有可选x值
matched = target_null.merge(
    source[['age', 'location', 'x']], 
    on=['age', 'location'], 
    how='left',
    suffixes=('', '_source')
)

# 3. 对每个原始行随机抽取1个匹配到的x值，无匹配的自动保留NaN
filled_x = matched.groupby('orig_idx')['x_source'].apply(lambda s: s.sample(n=1).iloc[0])

# 4. 填充结果回写原表
target.loc[filled_x.index, 'x'] = filled_x

方案2：预构建采样池（适合分组数远小于总行数的场景）

如果你的数据量在千万级以上，且age+location的分组数量远小于总行数，可以提前构建采样池进一步降低开销：

from collections import defaultdict
import random

# 预构建分组采样池，仅需执行1次
sample_pool = defaultdict(list)
for (age, loc), x_vals in source.groupby(['age', 'location'])['x']:
    sample_pool[(age, loc)] = x_vals.values

# 直接按行匹配采样
target['x'] = target.apply(
    lambda row: random.choice(sample_pool.get((row['age'], row['location']), [np.nan])) 
    if pd.isna(row['x']) else row['x'],
    axis=1
)

两种方案都完全符合填充规则：同组有放回随机采样，无匹配分组保留缺失值。

内容的提问来源于stack exchange，提问作者Artturi Björk