如何在Python Pandas中按两列分组填充三列NaN值且保留其他列数据

现有如下Pandas DataFrame：

import pandas as pd

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

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

需要实现以下需求：

• 按id和col_date（需转为datetime类型）分组
• 用分组内的有效值填充col1、col2、col3列的NaN值，每个分组至少有一行这三列有有效值
• col4、col5等其他列保持不变

期望得到的结果：

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

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

实现方法 #

通过Pandas的分组+填充逻辑即可完成，步骤如下：

1. 先将col_date转为datetime类型，确保分组逻辑准确：

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

2. 针对目标列执行分组填充：利用groupby按id和col_date分组，再通过transform对col1、col2、col3列填充NaN值。因为每个分组只有一组有效值，直接取分组内第一个非空值即可完成填充：

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

执行上述代码后，df就会变为期望的结果，col4、col5等其他列的数据完全保留不变。

完整代码示例 #

import pandas as pd

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

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

# 转换日期列类型
df['col_date'] = pd.to_datetime(df['col_date'])

# 填充指定列的NaN
fill_cols = ['col1', 'col2', 'col3']
df[fill_cols] = df.groupby(['id', 'col_date'])[fill_cols].transform(lambda x: x.fillna(x.first()))

print(df)

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