Pandas分组后按分位数映射列填充NaN的高效实现方案

问题背景 #

你有一个带分组列A、数值列B和含缺失值列C的DataFrame，已经按A分组算出了B的分位数排名B_q。现在需要给C里的每个NaN，用它所在组C的对应分位值（由B_q指定）来填充，但原来用apply的方法跑起来太慢，得换个高效的路子。

优化思路 #

别再逐行apply了，核心是先预计算每组的分位数映射表，再通过合并快速匹配填充，把循环开销降到最低。

步骤1：预生成每组的分位数-值映射 #

先按A分组，对每组非NaN的C计算对应B_q分位的数值，生成一个(A, B_q)到填充值的映射表：

import pandas as pd
import numpy as np

# 模拟示例数据（可替换为你的实际数据）
np.random.seed(42)
df = pd.DataFrame({
    'A': np.repeat(['X', 'Y', 'Z'], 1000),
    'B': np.random.randn(3000),
    'C': np.random.randn(3000)
})
# 随机生成C列的缺失值
df.loc[np.random.choice(df.index, 300), 'C'] = np.nan
# 按A分组计算B的分位数排名（示例用qcut分成10档，得到0-9的排名，对应你的B_q）
df['B_q'] = df.groupby('A')['B'].transform(lambda x: pd.qcut(x, 10, labels=False))

# 预计算每组的分位数映射：key是(A, B_q)，value是对应分位的C值
# 注意：如果你的B_q是0-1的连续分位值，直接用g['B_q'].unique()即可，无需除以10
quantile_map = df.groupby('A').apply(
    lambda g: g['C'].dropna().quantile(q=g['B_q'].unique() / 10)
).unstack().stack().reset_index(name='fill_val')
quantile_map.columns = ['A', 'B_q', 'fill_val']

步骤2：合并映射表快速填充 #

把原DataFrame和映射表按A、B_q合并，然后用fill_val填充C的NaN：

# 合并映射表
df = df.merge(quantile_map, on=['A', 'B_q'], how='left')
# 填充缺失值
df['C'] = df['C'].fillna(df['fill_val'])
# 删掉临时辅助列（可选操作）
df = df.drop('fill_val', axis=1)

为什么这个方法更快？ #

• 原来的apply是逐行处理，相当于给每一行重复执行一次分组分位数计算，冗余开销极大；
• 优化后仅给每组计算一次分位数映射，再通过Pandas的矢量化合并操作完成填充，时间复杂度从O(n)降到O(g*k + n)，其中g是组数，k是每组的分位数数量——数据量越大，性能提升越显著。

注意事项 #

• 如果你的B_q是0到1之间的连续分位值（比如0.25、0.5这类），直接把代码里的g['B_q'].unique() / 10改成g['B_q'].unique()即可；
• 若某组C的非NaN数据量过少，分位数计算可能不稳定，可添加判断逻辑，比如当非NaN数量低于阈值时，用组均值或预设默认值填充。

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