如何按user_id分组用c列均值原地填充a列NaN值

需求说明 #

数据集的a列存在NaN值，需按user_id分组，用每组c列的均值填充a列的NaN值，且需进行原地操作以适配大规模数据集。

初始数据集代码 #

import pandas as pd
import numpy as np

df = pd.DataFrame({'a': [0, np.nan, np.nan], 'user_id': [1, 2, 2], 'c': [3, 7, 7]})

print(df)

期望结果 #

df = pd.DataFrame({'a': [0, 7, 7], 'user_id': [1, 2, 2], 'c': [3, 7, 7]})
    
print(df)

尝试的代码及问题 #

执行以下代码后，a列仍存在NaN值：

df['a'].fillna(df.groupby('user_id')['a'].transform('mean'), inplace=True)

print(df)

问题原因 #

你之前的代码是计算每组a列的均值，但user_id=2的分组中a列全为NaN，计算出的均值也是NaN，填充后自然还是NaN。需求是用**c列的均值**来填充，而非a列自身的均值。

解决方案 #

修改分组计算的列为c，同时保持原地操作：

# 计算每组c列的均值，生成与原数据同长度的填充序列
fill_values = df.groupby('user_id')['c'].transform('mean')
# 原地填充a列的NaN值
df['a'].fillna(fill_values, inplace=True)

print(df)

执行后即可得到期望结果：

a  user_id  c
0  0        1  3
1  7        2  7
2  7        2  7

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