Pandas分组处理：替换指定列首次出现'A'前目标列值为NaN

示例数据 #

给定如下Pandas DataFrame，包含id（分组ID）、b（条件列）、c（目标列）：

import pandas as pd
import numpy as np

df = pd.DataFrame({'id' : [1,1,1,1,1,1,1,2,2,2,2,2,2,2], 
              'b'  : [3,4,5,'A',3,4,'A',1,'A',1,3,'A',2,3], 
              'c'  : [1,0,1,10,1,1,20,1,10,0,1,20,1,1]})
print(df)

输出：

id  b   c
0    1  3   1
1    1  4   0
2    1  5   1
3    1  A  10
4    1  3   1
5    1  4   1
6    1  A  20
7    2  1   1
8    2  A  10
9    2  1   0
10   2  3   1
11   2  A  20
12   2  2   1
13   2  3   1

需求 #

按id分组，将每个分组中b列首次出现'A'之前的c列值替换为np.nan，期望输出如下：

desired_output_df = pd.DataFrame({'id' : [1,1,1,1,1,1,1,2,2,2,2,2,2,2], 
                                  'b'  : [3,4,5,'A',3,4,'A',1,'A',1,3,'A',2,3], 
                                  'c'  : [np.nan,np.nan,np.nan,10,1,1,20,np.nan,10,0,1,20,1,1]})
print(desired_output_df)

输出：

id  b     c
0    1  3   NaN
1    1  4   NaN
2    1  5   NaN
3    1  A  10.0
4    1  3   1.0
5    1  4   1.0
6    1  A  20.0
7    2  1   NaN
8    2  A  10.0
9    2  1   0.0
10   2  3   1.0
11   2  A  20.0
12   2  2   1.0
13   2  3   1.0

遇到的问题 #

通过代码df.groupby('id').apply(lambda x: x.loc[:(x.b == 'A').idxmax()-1]).index获取到需要修改的索引，但结果为MultiIndex，无法直接用于替换值。

解决方案 #

方法1：分组内直接标记替换 #

无需提取索引，直接在分组内生成标记逻辑处理：

def process_group(group):
    # 获取分组内首次出现'A'的位置
    first_A_pos = (group['b'] == 'A').idxmax()
    # 标记首次'A'之前的行
    mask = group.index < first_A_pos
    # 替换对应位置的c值为NaN
    group['c'] = np.where(mask, np.nan, group['c'])
    return group

df_result = df.groupby('id').apply(process_group).reset_index(drop=True)
print(df_result)

方法2：提取MultiIndex中的原始索引 #

如果要保留提取索引的思路，可从MultiIndex中取出原DataFrame的行索引再赋值：

# 提取需要修改的原始行索引，用explode展开MultiIndex
target_idx = df.groupby('id').apply(lambda x: x.loc[:(x.b == 'A').idxmax()-1].index).explode().values

# 直接替换对应位置的c列值
df.loc[target_idx, 'c'] = np.nan
print(df)

两种方法均可得到符合需求的结果，推荐第一种，逻辑更直观且避免索引操作的潜在问题。

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