按ID分组结合length规则填充DataFrame中的NaN值

输入数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({"ID" : [1, 1, 1, 2, 2, 2, 2],
                  "length" : [0.7, 0.7, 0.7, 0.8, 0.6, 0.6, 0.7],
                  "height" : [7, 9, np.nan, 4, 8, np.nan, 5]})

生成的DataFrame如下：

ID  length  height
0   1   0.7     7
1   1   0.7     9
2   1   0.7     NaN
3   2   0.8     4
4   2   0.6     8
5   2   0.6     NaN
6   2   0.7     5

填充规则 #

按ID分组后执行以下填充逻辑：

• 若分组内所有length值相同，用该组height的最大值填充NaN；
• 若分组内length值不唯一，用该组中最大length对应的height值填充NaN。

期望输出 #

ID  length  height
0   1   0.7     7
1   1   0.7     9
2   1   0.7     9
3   2   0.8     4
4   2   0.6     8
5   2   0.6     4
6   2   0.7     5

实现代码 #

# 按ID分组处理的自定义函数
def fill_nan(group):
    # 判断分组内length是否全部一致
    if group['length'].nunique() == 1:
        fill_val = group['height'].max()
    else:
        # 找到分组内最大的length值
        max_length = group['length'].max()
        # 提取该length对应的height值作为填充值
        fill_val = group[group['length'] == max_length]['height'].iloc[0]
    # 填充当前分组的NaN值
    group['height'] = group['height'].fillna(fill_val)
    return group

# 应用分组逻辑并重置索引
df_filled = df.groupby('ID').apply(fill_nan).reset_index(drop=True)
print(df_filled)

代码说明 #

1. 分组判断逻辑：通过nunique()快速判断分组内length是否唯一，以此选择不同的填充值来源；
2. 填充值提取：当length不唯一时，先定位最大length，再提取对应的height值；
3. 分组应用：利用groupby().apply()将填充逻辑批量应用到每个ID分组，最后重置索引得到规整的结果。

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