按组查找首个指定值并为组内后续行填充该值的Pandas实现

问题描述 #

现有如下Pandas DataFrame：

import pandas as pd
import numpy as np

data = {'Name': ['Tom', 'Tom', 'Tom', 'Tom'],
        'Sequence': [2021, 2022, 2023, 2024],
        'Car': ['Ford', 'Chevy', 'Tesla', np.nan]}
df = pd.DataFrame(data)

需要创建new_car列：按Name分组后，找到组内Car列第一个出现'Chevy'的位置，将该位置及之后的所有行new_car设为'Chevy'，之前的行设为NaN，预期结果为[NaN, 'Chevy', 'Chevy', 'Chevy']。

方法一：分组自定义函数处理 #

通过自定义函数对每个分组单独处理，精准定位第一个'Chevy'的位置：

def process_group(group):
    # 获取组内第一个'Chevy'的索引
    first_chevy_idx = group[group['Car'] == 'Chevy'].index.min()
    # 处理组内无'Chevy'的情况
    if pd.isna(first_chevy_idx):
        group['new_car'] = np.nan
    else:
        # 索引大于等于目标位置的行设为'Chevy'，其余为NaN
        group['new_car'] = np.where(group.index >= first_chevy_idx, 'Chevy', np.nan)
    return group

# 分组应用函数，group_keys=False避免额外添加分组键索引
df = df.groupby('Name', group_keys=False).apply(process_group)

方法二：标记+累积和简化实现 #

利用累积和的特性，无需手动定位索引，更简洁高效：

# 标记'Chevy'出现的位置为1，其余为0
df['chevy_flag'] = (df['Car'] == 'Chevy').astype(int)
# 按Name分组计算累积和：第一个'Chevy'之后的行累积和始终≥1
df['new_car'] = df.groupby('Name')['chevy_flag'].transform(lambda x: x.cumsum())
# 根据累积和结果赋值，≥1的设为'Chevy'，否则为NaN
df['new_car'] = df['new_car'].map(lambda x: 'Chevy' if x >= 1 else np.nan)
# 删除中间辅助列
df = df.drop('chevy_flag', axis=1)

验证结果 #

执行后df的new_car列将符合预期：

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