如何在Python Pandas中按两列分组填充DataFrame三列的NaN值

需求说明 #

按id和col_date（需转为datetime类型）分组，用每组内col1、col2、col3的有效值填充同组内对应列的缺失值，且每组至少存在一行有效值。

实现步骤 #

1. 导入依赖并构造原始DataFrame

import pandas as pd

data = [
    (1, None, None, None, '2023-01-10'),
    (1, None, None, None, '2023-01-10'),
    (1, 9, 0, 0.55, '2023-01-10'),
    (2, None, None, None, '2023-11-22'),
    (2, 88, 1, 0.68, '2023-11-22')
]

df = pd.DataFrame(data, columns=['id', 'col1', 'col2', 'col3', 'col_date'])

2. 转换col_date为datetime类型

df['col_date'] = pd.to_datetime(df['col_date'])

3. 分组填充缺失值
   提供两种可行方法：

方法一：取每组非空值填充

利用transform提取每组的第一个非空值，填充组内所有缺失：

fill_cols = ['col1', 'col2', 'col3']
df[fill_cols] = df.groupby(['id', 'col_date'])[fill_cols].transform(
    lambda x: x.fillna(x.dropna().iloc[0])
)

方法二：双向填充（更通用）

先向后填充（用后续行的有效值补前面缺失），再向前填充（用前面行的有效值补后面缺失），适配有效值在组内任意位置的场景：

fill_cols = ['col1', 'col2', 'col3']
df[fill_cols] = df.groupby(['id', 'col_date'])[fill_cols].transform(lambda x: x.bfill().ffill())

验证结果 #

执行上述代码后，输出的DataFrame即为需求中的目标结果：

print(df)
# 输出：
#    id  col1  col2  col3   col_date
# 0   1     9     0  0.55 2023-01-10
# 1   1     9     0  0.55 2023-01-10
# 2   1     9     0  0.55 2023-01-10
# 3   2    88     1  0.68 2023-11-22
# 4   2    88     1  0.68 2023-11-22

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