Pandas分组处理DataFrame：按最大值条件生成目标列需求

问题背景 #

现有如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({'c1': [1,1,1,1,1,-1,-1,-1], 'c2': [0.1, 0.02, 0.5,0.07,0.09,0.8,0.03,0.02]} )

已通过以下代码添加了c3列，该列为c1分组后对应组内c2的最大值：

df['c3'] = df.groupby('c1')['c2'].transform('max')

需求：针对c1的每个分组，若组内c2的最大值≥0.5，则将该分组中最大值所在行之后的所有行的目标列设为0，最大值所在行及之前的行保留c1的值；若最大值<0.5，则目标列保留c1的值。最终目标列效果如下：

解决方案 #

可以通过分组生成掩码标记最大值位置，再结合条件判断生成目标列，代码如下：

完整代码 #

import pandas as pd
import numpy as np

# 初始化数据
df = pd.DataFrame({'c1': [1,1,1,1,1,-1,-1,-1], 'c2': [0.1, 0.02, 0.5,0.07,0.09,0.8,0.03,0.02]} )
# 添加组内最大值列c3
df['c3'] = df.groupby('c1')['c2'].transform('max')

# 生成掩码：标记每个分组内当前行是否在最大值行及之前
max_position_mask = df.groupby('c1')['c2'].transform(lambda x: x.cummax() == x.max())

# 根据条件生成目标列
df['desired_col'] = np.where(
    df['c3'] >= 0.5,
    np.where(max_position_mask, df['c1'], 0),
    df['c1']
)

print(df)

代码逻辑解释 #

1. 生成最大值位置掩码：x.cummax()会计算分组内c2的累积最大值，当累积最大值等于组内最大值x.max()时，说明当前行在最大值所在行或之前，以此生成布尔掩码。
2. 条件判断生成目标列：
   • 第一层判断：如果组内最大值c3≥0.5，进入第二层判断；否则直接保留c1的值。
   • 第二层判断：若掩码为True（在最大值行及之前），则取c1的值；否则设为0。

运行代码后即可得到符合需求的desired_col列。

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